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库来正确管理和终止子进程。这种方法能够确保系统资源的合理使用,并避免出现僵尸进程问题。希望本文对你在实际项目中处理类似问题时有所帮助。

相关推荐
码事漫谈4 分钟前
AI行业热点抓取和排序系统实现案例
后端
我要学习别拦我~11 分钟前
kaggle分析项目:steam付费游戏数据分析
python·游戏·数据分析
方圆想当图灵12 分钟前
关于 Nacos 在 war 包部署应用关闭部分资源未释放的原因分析
后端
大模型真好玩17 分钟前
深入浅出LangChain AI Agent智能体开发教程(四)—LangChain记忆存储与多轮对话机器人搭建
前端·人工智能·python
love530love19 分钟前
命令行创建 UV 环境及本地化实战演示—— 基于《Python 多版本与开发环境治理架构设计》的最佳实践
开发语言·人工智能·windows·python·conda·uv
Lemon程序馆23 分钟前
今天聊聊 Mysql 的那些“锁”事!
后端·mysql
龙卷风040525 分钟前
使用本地IDEA连接服务器远程构建部署Docker服务
后端·docker
vv安的浅唱30 分钟前
Golang基础笔记七之指针,值类型和引用类型
后端·go
陪我一起学编程41 分钟前
MySQL创建普通用户并为其分配相关权限的操作步骤
开发语言·数据库·后端·mysql·oracle