为循环添加Multiprocessing后计算耗时变长？千万级数据集优化求助

一、代码中的致命错误 #

• 循环变量误用：calc函数的参数是chunk（单块数据），但循环里写的是chunks.itertuples（全部数据块），这会导致每个进程都重复处理所有客户数据，相当于4个进程做了4遍单进程的工作，耗时必然大幅增加。
  修正：将chunks.itertuples改为chunk.itertuples。

二、多进程开销优化点 #

1. 减少大对象的跨进程复制 #

matrix_data、lookup1、end_dates这类全局变量会被每个子进程完整复制一份，当数据量较大时，会产生巨大的内存开销和序列化时间。优化方式：

• 使用multiprocessing.Manager共享只读数据，或用multiprocessing.shared_memory存储numpy数组，避免内存重复占用；
• 直接在子进程内部加载这些静态数据，跳过父进程到子进程的序列化传递步骤。

2. 合理设置进程数量 #

不要固定设置num_processes=4，应根据Sagemaker实例的CPU核心数动态调整，一般设为os.cpu_count()或os.cpu_count()-1，避免进程过多导致CPU上下文切换频繁。示例：

import os
num_processes = os.cpu_count()

3. 降低进程间数据传递开销 #

用mp.Queue传递百万级结果字典会产生巨大的序列化/反序列化开销，建议改用：

• multiprocessing.Pool的map/imap方法，自动处理结果收集，比手动维护Queue更高效；
• 按chunk将结果写入磁盘（如parquet格式），最后再合并，减少内存占用和进程间数据传递。

三、计算逻辑本身的优化 #

1. 向量化替代Python循环 #

当前逐客户的循环可以用numpy向量化操作重构，大幅减少Python解释器的循环开销。

2. 预计算重复值 #

提前提取固定值，避免每次循环重复计算或查询：

# 预计算lookup表的值和日期归一化结果
lookup_values = lookup1['lim %'].values
end_dates_normalized = end_dates / 365

四、Sagemaker Studio环境限制排查 #

• 检查实例资源：如果实例内存不足，多进程运行时会触发内存交换（swap），导致性能骤降。可通过top命令查看资源占用，必要时升级实例类型；
• 确认CPU亲和性：部分Sagemaker容器可能对进程CPU亲和性有限制，导致进程无法充分利用核心，可通过taskset命令查看或调整。

修正后的示例代码 #

import os
import numpy as np
import pandas as pd
import multiprocessing as mp

# 预计算重复值
lookup_values = lookup1['lim %'].values
end_dates_normalized = end_dates / 365

def calc(chunk):
    results1 = {}
    # 修正循环变量：使用当前进程分配的chunk
    for cust_id, state, amount, start, group, loan_rate in chunk.itertuples(name=None, index=False):
        res1 = [amount * matrix_data[start-1, state, :]]
        for year in range(start+1, len(matrix_data)+1):
            res1.append(lookup_values[year-1] * np.array(res1[-1]))
            res1.append(res1[-1] * loan_rate * end_dates_normalized[year-1])
            res1.append(res1[-1] + 100)
            res1.append(np.linalg.multi_dot([res1[-1], matrix_data[year-1]]))
        results1[cust_id] = res1
    return results1

if __name__ == '__main__':
    num_processes = os.cpu_count()
    chunks = np.array_split(data1, num_processes)
    
    # 使用Pool简化进程管理与结果收集
    with mp.Pool(num_processes) as pool:
        chunk_results = pool.map(calc, chunks)
    
    # 合并所有chunk的结果
    results1 = {}
    for res in chunk_results:
        results1.update(res)

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