Python multiprocess Pool.map初始化耗时过长的优化咨询

原代码 #

%%time
import pandas as pd
import numpy as np  # 补充原代码缺失的导入
from datetime import datetime
from multiprocessing import Pool
from functools import partial  # 补充原代码缺失的导入


def get_id(x):
    return x['motor_id']

def parallel_load(df, get_id):
    print('parallel_load: ', datetime.now())
    return df.apply(lambda x: get_id(x), axis=1)

def parallel_df(df, func, n_cores):
    print('start: ', datetime.now())
    df_split = np.array_split(df, n_cores)
    print('splitted: ', datetime.now())
    pool = Pool(n_cores)
#     df = pd.concat(pool.map(func, df_split))
    s = datetime.now()
    print('pool start: {}'.format(s))
    data = pool.map(func, df_split)
    print('pool end: {}, {}s'.format(datetime.now(), (datetime.now() - s).seconds))
    pool.close()
    pool.join()
    print('pool join: {}s'.format((datetime.now() - s).seconds))
    
    return data

new_func = partial(
    parallel_load, get_id=get_id
)

ndf = pd.DataFrame({'motor_id': np.arange(10000)})
parallel_df(ndf, new_func, 4)

运行结果 #

start:  2024-01-02 16:58:09.390751
splitted:  2024-01-02 16:58:09.391897
parallel_load:  parallel_load: parallel_load: parallel_load: parallel_load:  parallel_load: 2024-01-02 16:58:12.453533parallel_load:  parallel_load: 2024-01-02 16:58:12.456901parallel_load: parallel_load:   

2024-01-02 16:58:12.456053  2024-01-02 16:58:12.4578302024-01-02 16:58:12.457674   2024-01-02 16:58:12.457189

2024-01-02 16:58:12.457663

2024-01-02 16:58:12.4583492024-01-02 16:58:12.458657
2024-01-02 16:58:12.457537


pool start: 2024-01-02 16:58:12.448087
pool end: 2024-01-02 16:58:12.504469, 0s
pool join: 0s
CPU times: user 52.5 ms, sys: 3.48 s, total: 3.53 s
Wall time: 3.39 s

优化方案 #

小任务场景下，multiprocessing.Pool的进程启动成本（进程创建、内存复制、模块导入等）远高于任务本身的处理时间，可通过以下方式优化：

1. 改用线程池替代进程池

线程共享同一进程的内存空间，无需重复导入模块和复制数据，启动成本远低于进程池，适合无GIL阻塞的简单计算或IO密集型任务。

%%time
import pandas as pd
import numpy as np
from datetime import datetime
from multiprocessing.dummy import Pool  # 替换为线程池
from functools import partial


def get_id(x):
    return x['motor_id']

def parallel_load(df, get_id):
    print('parallel_load: ', datetime.now())
    return df.apply(lambda x: get_id(x), axis=1)

def parallel_df(df, func, n_cores):
    print('start: ', datetime.now())
    df_split = np.array_split(df, n_cores)
    print('splitted: ', datetime.now())
    pool = Pool(n_cores)
    s = datetime.now()
    print('pool start: {}'.format(s))
    data = pool.map(func, df_split)
    print('pool end: {}, {}s'.format(datetime.now(), (datetime.now() - s).seconds))
    pool.close()
    pool.join()
    print('pool join: {}s'.format((datetime.now() - s).seconds))
    
    return data

new_func = partial(
    parallel_load, get_id=get_id
)

ndf = pd.DataFrame({'motor_id': np.arange(10000)})
parallel_df(ndf, new_func, 4)

2. 复用进程池

若需多次执行同类小任务，提前初始化全局进程池，将启动成本平摊到多个任务上，避免重复创建进程。

# 提前初始化全局进程池
import pandas as pd
import numpy as np
from datetime import datetime
from multiprocessing import Pool
from functools import partial

global_pool = Pool(4)

def get_id(x):
    return x['motor_id']

def parallel_load(df, get_id):
    print('parallel_load: ', datetime.now())
    return df.apply(lambda x: get_id(x), axis=1)

def parallel_df(df, func):
    print('start: ', datetime.now())
    df_split = np.array_split(df, 4)
    print('splitted: ', datetime.now())
    s = datetime.now()
    print('pool start: {}'.format(s))
    data = global_pool.map(func, df_split)
    print('pool end: {}, {}s'.format(datetime.now(), (datetime.now() - s).seconds))
    print('pool join: {}s'.format((datetime.now() - s).seconds))
    return data

new_func = partial(
    parallel_load, get_id=get_id
)

# 第一次执行
ndf = pd.DataFrame({'motor_id': np.arange(10000)})
parallel_df(ndf, new_func)

# 第二次执行，复用已初始化的进程池
ndf2 = pd.DataFrame({'motor_id': np.arange(10000,20000)})
parallel_df(ndf2, new_func)

# 任务完成后关闭进程池
global_pool.close()
global_pool.join()

3. 优化任务粒度

减少进程数量或合并任务，让每个进程处理更大的数据集，抵消进程启动成本。比如本例中10000条数据可仅拆分为2份，甚至单进程处理。

4. 简化任务逻辑

避免不必要的计算开销，比如将df.apply替换为更高效的直接列提取，缩短任务本身的处理时间：

def parallel_load(df, get_id):
    print('parallel_load: ', datetime.now())
    return df['motor_id'].tolist()  # 直接提取列数据，比apply更高效

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