Flask使用Celery与多进程管理:优雅处理长时间任务与子进程终止技巧(multiprocessing)(subprocess)

在许多任务处理系统中,我们需要使用异步任务队列来处理繁重的计算或长时间运行的任务,如模型训练。Celery是一个广泛使用的分布式任务队列,而在某些任务中,尤其是涉及到调用独立脚本的场景中,我们需要混合使用multiprocessingsubprocess模块来启动和管理这些任务进程。然而,这种组合有时会带来一些挑战,如进程冲突和子进程无法正确终止的问题。

本文将讨论如何使用Celery、Multiprocessing和Subprocess来处理这些问题,并在需要时正确关闭子进程,实现完美的进程管理与切换。

问题描述

当通过celery.control.revoke来终止Celery任务时,如果任务启动了多个子进程,例如使用multiprocessingsubprocess模块,这些子进程不会被立即终止。在某些情况下,子进程会继续运行,导致任务无法彻底停止,并可能造成系统资源浪费。

解决方案

我们可以通过组合使用psutil库来实现对子进程的监控和终止,从而确保所有相关的进程都能正确关闭。以下是具体实现步骤。

安装必要的库

确保你已经安装psutil库:

复制代码
pip install psutil

修改代码实现

1. Celery任务与Multiprocessing结合Subprocess

首先,我们创建一个Celery任务。当任务启动时,它会使用multiprocessing模块启动一个新的进程,该进程将执行独立的Python脚本。代码如下:

python 复制代码
import json
import os
import psutil
import multiprocessing
import subprocess
from celery_app import celery

import torch.multiprocessing as mp

mp.set_start_method('spawn', True)

def run_script(json_test_path, uid):
    command = f"python training.py {json_test_path}"
    process = subprocess.Popen(command, shell=True)
    print("===========================================PID:", process.pid)
    print("===========================================uid:", uid)

    process.wait()
    return process.pid

@celery.task(bind=True)
def lora_train_task(self, json_test_demo):
    # 将json_test保存到临时文件中
    json_test_path = f"training_config_{json_test_demo['uid']}.json"
    json_test_path = os.path.abspath(json_test_path)
    with open(json_test_path, 'w') as f:
        json.dump(json_test_demo, f)

    # 使用多进程调用独立脚本
    p = multiprocessing.Process(target=run_script, args=(json_test_path, json_test_demo['uid']))
    p.start()
    p.join()

    return 0
2. 使用psutil关闭子进程

我们通过调用psutil库来监控并关闭所有相关的子进程。以下是实现终止任务和子进程的代码示例:

python 复制代码
import psutil

def terminate_process_tree(pid):
    try:
        parent = psutil.Process(pid)
        for child in parent.children(recursive=True):  # This will recursively find all child processes
            child.terminate()
        parent.terminate()
    except psutil.NoSuchProcess:
        pass

# 示例:终止任务时调用终止子进程函数
celery.control.revoke(args['task_id'], terminate=True, signal='SIGKILL')
if args['pid_id']:
    terminate_process_tree(int(args['pid_id']))

完整示例

将上述代码组合起来,我们得到完成的实现。如下所示:

python 复制代码
import json
import os
import psutil
import multiprocessing
import subprocess
from celery_app import celery
import torch.multiprocessing as mp

mp.set_start_method('spawn', True)

def run_script(json_test_path, uid):
    command = f"python training.py {json_test_path}"
    process = subprocess.Popen(command, shell=True)
    print("===========================================PID:", process.pid)
    print("===========================================uid:", uid)

    process.wait()
    return process.pid

@celery.task(bind=True)
def lora_train_task(self, json_test_demo):
    # 将json_test保存到临时文件中
    json_test_path = f"training_config_{json_test_demo['uid']}.json"
    json_test_path = os.path.abspath(json_test_path)
    with open(json_test_path, 'w') as f:
        json.dump(json_test_demo, f)

    # 使用多进程调用独立脚本
    process = multiprocessing.Process(target=run_script, args=(json_test_path, json_test_demo['uid']))
    process.start()
    process.join()

    return 0

def terminate_process_tree(pid):
    try:
        parent = psutil.Process(pid)
        for child in parent.children(recursive=True):  # 递归找到所有子进程
            child.terminate()
        parent.terminate()
    except psutil.NoSuchProcess:
        pass

# 示例:终止任务时调用终止子进程函数
celery.control.revoke(args['task_id'], terminate=True, signal='SIGKILL')
if args['pid_id']:
    terminate_process_tree(int(args['pid_id']))

结论

本文示范了如何通过混合使用Celery、Multiprocessing与Subprocess来处理复杂的任务执行场景,同时介绍了通过psutil库来正确管理和终止子进程。这种方法能够确保系统资源的合理使用,并避免出现僵尸进程问题。希望本文对你在实际项目中处理类似问题时有所帮助。

相关推荐
YQ_ZJH6 分钟前
Spring Boot 如何校验前端传递的参数
前端·spring boot·后端
普罗米修斯13 分钟前
C++ 设计模式理论与实战大全【共73课时】
c++·后端
普罗米修斯16 分钟前
C++ 设计模式原理与实战大全-架构师必学课程 | 完结
c++·后端
MX_935930 分钟前
SpringBoot项目优先级以及bean的管理
java·spring boot·后端·spring
豆浆Whisky1 小时前
Go泛型实战指南:从入门到工程最佳实践|Go语言进阶(12)
后端·go
可触的未来,发芽的智生1 小时前
触摸未来2025.10.04:当神经网络拥有了内在记忆……
人工智能·python·神经网络·算法·架构
元亓亓亓1 小时前
SSM--day4--SpringMVC(补充)
java·后端·ssm
蔗理苦1 小时前
2025-10-07 Python不基础 20——全局变量与自由变量
开发语言·python
xiaohanbao092 小时前
理解神经网络流程
python·神经网络
韩立学长2 小时前
【开题答辩实录分享】以《基于Python的旅游网站数据爬虫研究》为例进行答辩实录分享
python·旅游