如何指定Python程序运行核心数并监控其内存占用与执行时间

原有问题修复 #

• p.map运行报错问题：multiprocessing.Pool.map() 要求传入两个必填参数：第一个是执行函数，第二个是可迭代对象，可迭代对象的每个元素会依次作为参数传给执行函数。所以你需要给process函数加一个占位参数（哪怕你业务逻辑用不上这个参数），同时根据你实际要跑的任务数量传入对应长度的可迭代对象（比如要跑100个任务就传range(100)）。
• 内存测量不准确问题：原有代码取的是系统总内存大小，完全不是你程序的实际占用。而且进程池运行结束后子进程会自动销毁，只在结束时刻采样也拿不到运行峰值。正确的做法是在任务运行过程中定时采样进程池所有子进程的物理内存占用，记录峰值作为该核心数下的内存指标。

完整可运行代码 #

import time
import psutil
import pandas as pd
from multiprocessing import Pool
from threading import Thread

# 可配置参数
TASK_COUNT = 200  # 实际要并行处理的任务总数，按自己业务改
SAMPLING_INTERVAL = 0.1  # 内存采样间隔，单位秒
test_cores = [2,4,6,8,10,12,14,16]  # 要测试的核心数范围

def process(task_id):
    # 替换成你自己的业务逻辑即可，task_id是map传入的占位参数，不需要可以忽略
    res = 0
    for i in range(1000000):
        res += i
    return res

def memory_monitor(peak_memory):
    # 后台内存采样线程逻辑
    while True:
        # 获取当前主进程的所有子进程（即进程池的工作进程）
        children = psutil.Process().children()
        total_rss = 0
        for child in children:
            try:
                total_rss += child.memory_info().rss
            except psutil.NoSuchProcess:
                # 子进程已经退出，跳过
                continue
        # 换算成MB并更新峰值
        current_mb = total_rss / 1024 / 1024
        if current_mb > peak_memory[0]:
            peak_memory[0] = current_mb
        time.sleep(SAMPLING_INTERVAL)
        # 所有子进程退出则结束采样
        if not children:
            break

if __name__ == '__main__':
    result = {'num_cores': [], 'peak_memory_mb': [], 'execution_time_s': []}
    for cores in test_cores:
        print(f"正在测试{cores}核运行情况...")
        peak_memory = [0]  # 用列表存方便线程内修改
        # 启动内存监控线程
        monitor_thread = Thread(target=memory_monitor, args=(peak_memory,))
        monitor_thread.start()
        
        start_time = time.time()
        with Pool(cores) as pool:
            # 传入任务可迭代对象
            pool.map(process, range(TASK_COUNT))
        end_time = time.time()
        
        monitor_thread.join()
        # 记录结果
        result['num_cores'].append(cores)
        result['peak_memory_mb'].append(round(peak_memory[0], 2))
        result['execution_time_s'].append(round(end_time - start_time, 2))
    
    # 导出结果
    df = pd.DataFrame(result)
    print(df)
    # df.to_csv("core_test_result.csv", index=False) # 可选保存到csv

核心逻辑说明 #

• 内存采样用了独立的后台线程，每隔100ms采集一次所有子进程的RSS内存（实际物理内存占用，单位MB），取运行全程的最大值作为统计结果，避免漏记峰值
• 你可以根据自己的业务需求修改TASK_COUNT（要并行处理的任务总数）、SAMPLING_INTERVAL（内存采样间隔，单位秒）、test_cores（要测试的核心数范围）三个参数
• 最终生成的DataFrame会包含所有测试核心数对应的执行时间和内存峰值，你可以直接基于这份数据做权衡分析：通常核心数超过物理CPU核心数之后，执行时间下降的边际效益会快速衰减，同时内存占用还会持续上升，你可以根据自己对耗时和内存的约束选最优参数

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