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

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

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

问题描述

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

解决方案

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

安装必要的库

确保你已经安装psutil库：

pip install psutil

修改代码实现

1. Celery任务与Multiprocessing结合Subprocess

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

import json
import os
import psutil
import multiprocessing
import subprocess
from celery_app import celeryimport torch.multiprocessing as mpmp.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库来监控并关闭所有相关的子进程。以下是实现终止任务和子进程的代码示例：

import psutildef terminate_process_tree(pid):try:parent = psutil.Process(pid)for child in parent.children(recursive=True):  # This will recursively find all child processeschild.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']))

完整示例

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

import json
import os
import psutil
import multiprocessing
import subprocess
from celery_app import celery
import torch.multiprocessing as mpmp.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 0def 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']))