Python使用Pool多线程时内存超阈值如何强制终止所有进程

解决方案 #

注意：你代码中使用的multiprocessing.Pool属于多进程模型而非多线程，进程之间内存空间独立，要终止所有工作进程可以通过「子进程触发阈值时抛异常→主进程捕获异常后终止整个进程池」的逻辑实现。

修正后的完整代码 #

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

# 自定义内存超限异常，用于区分普通报错和内存触发的终止逻辑
class MemoryOverLimitError(Exception):
    pass

def my_func(item):
    a = MemoryConsumingObject()
    # 修正内存读取逻辑：获取已使用内存数值，转整数后再比较
    memory_usage_giga = int(os.popen('free -g').readlines()[1].split()[2])
    if memory_usage_giga > 50:
        raise MemoryOverLimitError("内存占用超过50G阈值")
    return a

if __name__ == "__main__":
    my_list = list(np.arange(1,1000))
    pool = Pool(cpu_count())
    try:
        # 迭代imap结果触发任务执行
        result = list(pool.imap(my_func, my_list))
    except MemoryOverLimitError:
        # 捕获内存超限异常，立刻终止所有子进程
        pool.terminate()
        pool.join()
        print("已触发内存阈值，所有工作进程已终止")
    else:
        # 所有任务正常完成的收尾逻辑
        pool.close()
        pool.join()

关键逻辑说明 #

• 不要在子进程内直接操作进程池，子进程没有权限控制父进程创建的池资源，抛异常是最稳妥的跨进程通知方式
• pool.terminate()会立刻停止所有未完成的工作进程，不会等待任务执行完毕，完全符合内存超限紧急终止的需求
• 原代码内存读取的字段索引有误：free -g输出第二行的第三个字段才是已使用内存，之前的索引取到的是总内存，会导致阈值判断完全失效

Jupyter环境适配说明 #

Jupyter环境中运行多进程代码不需要额外修改，终止进程池后不会残留僵尸进程，内存会被系统自动回收。如果需要终止后立刻释放当前Python进程持有的内存，可以在终止逻辑后加del a, result手动触发垃圾回收。

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