如何在Python子进程中稳健地捕获PyTorch的 GPU内存不足（OOM）异常，并在一个自主AI代理循环中动态调整批量大小？

人工智能 2026-07-11

我正在构建一个自治AI代理（负责管理训练工作流），它会自动生成PyTorch/OpenMMLab的训练脚本，并在后台子进程中执行它们。

代理常见的问题之一是 RuntimeError: CUDA out of memory。目前，代理从 subprocess.Popen 的标准错误输出读取，解析堆栈跟踪，并重写配置文件以将 batch_size 减半。

然而，子进程有时会在GPU上留下僵死的CUDA进程或未释放的内存，导致后续的自动重试直接失败。

当前简化的做法：

import subprocess
import re

def run_training_agent.py(config_path):
    process = subprocess.Popen(["python", "train.py", config_path], stderr=subprocess.PIPE)
    _, stderr = process.communicate()

    if "CUDA out of memory" in stderr.decode('utf-8'):
        # Agent logic to rewrite batch_size in config_path
        reduce_batch_size(config_path) 
        # Restart
        run_training_agent(config_path)

我的问题：

最稳健的方式是什么，能够确保原始子进程在代理启动重试之前完全释放VRAM？我应该使用Python的 multiprocessing 代替 subprocess 吗？
是否有办法通过钩子在PyTorch脚本内部拦截OOM异常，并把确切的内存需求反馈给父进程的代理，而不是仅仅解析原始的 stderr 文本？

对于构建健壮的自动化ML流水线，若有见解，将不胜感激！

解决方案

我想给出一个答案。

import subprocess
import os
import signal

def run_training_agent(config_path):
    process = subprocess.Popen(
        ["python", "train.py", config_path],
        stderr=subprocess.PIPE,
        preexec_fn=os.setsid  # create new process group
    )

    _, stderr = process.communicate()

    if "CUDA out of memory" in stderr.decode():
        # Kill EVERYTHING in the group
        os.killpg(os.getpgid(process.pid), signal.SIGTERM)

        reduce_batch_size(config_path)

        run_training_agent(config_path)

PyTorch