Python中Celery任务被撤销时如何关闭Selenium WebDriver

你遇到的这个问题很常见——app.control.revoke()只是终止了Celery Worker进程，但WebDriver是一个独立的子进程/浏览器实例，Celery的撤销操作不会自动清理它。下面是几个靠谱的解决思路，按实用性排序：

方案1：利用Celery的task_revoked信号做清理 #

Celery提供了任务生命周期的信号，其中task_revoked会在任务被撤销时触发。你可以在这个信号的处理函数里，针对被撤销的任务ID，找到对应的WebDriver实例并关闭它。

首先，你需要在任务运行时，把任务ID和对应的WebDriver实例关联存储起来（单Worker场景用全局字典即可，多Worker建议用Redis这类共享存储）：

# 比如在你的任务模块里
from celery import signals
from selenium import webdriver

# 单Worker场景用全局字典，多Worker换Redis/数据库
driver_map = {}

@signals.task_revoked.connect
def handle_task_revoked(sender=None, request=None, **kwargs):
    task_id = request.id
    if task_id in driver_map:
        driver = driver_map[task_id]
        try:
            driver.quit()  # 用quit()而非close()，quit会彻底终止驱动进程
            del driver_map[task_id]
            print(f"已为任务{task_id}关闭WebDriver")
        except Exception as e:
            print(f"关闭WebDriver失败: {str(e)}")

# 你的Celery任务
@app.task(bind=True)
def run_selenium_task(self):
    driver = webdriver.Chrome()
    # 绑定当前任务ID与Driver实例
    driver_map[self.request.id] = driver
    try:
        # 这里写你的Selenium业务逻辑
        driver.get("https://example.com")
        # ... 其他操作
    finally:
        # 任务正常完成时也要清理字典和Driver
        driver.quit()
        if self.request.id in driver_map:
            del driver_map[self.request.id]

如果是多Worker分布式场景，全局字典不生效，这时候可以存储任务ID和Driver的进程ID，在信号处理函数里调用系统命令（比如pkill/taskkill）杀死对应的浏览器进程。

方案2：在任务内部捕获终止信号 #

当你用terminate=True调用revoke时，Celery会给Worker进程发送SIGTERM信号（Windows下是类似的终止信号）。你可以在任务里注册信号处理函数，收到终止信号时主动关闭WebDriver。

示例代码：

import signal
from selenium import webdriver
from celery import Celery

app = Celery('proj', broker='redis://localhost:6379/0')

def handle_termination(signum, frame, driver):
    print("收到终止信号，正在关闭WebDriver...")
    driver.quit()
    exit(0)

@app.task(bind=True)
def selenium_task(self):
    driver = webdriver.Chrome()
    # 注册SIGTERM信号处理函数，传入driver实例
    signal.signal(signal.SIGTERM, lambda s, f: handle_termination(s, f, driver))
    
    try:
        # 业务逻辑
        driver.get("https://example.com")
        # ... 耗时操作
    except Exception as e:
        driver.quit()
        raise e

注意：Windows系统需要调整信号类型（比如signal.CTRL_BREAK_EVENT），如果Worker用了solo以外的进程池（比如prefork），信号处理逻辑可能需要额外适配。

方案3：用上下文管理器做兜底 #

不管任务是正常完成还是被终止，用上下文管理器（with语句）管理WebDriver的生命周期，确保资源被释放。虽然强制终止时finally块可能不执行，但上下文的__exit__方法多数情况下会触发。

示例：

from selenium import webdriver
from celery import Celery

app = Celery('proj', broker='redis://localhost:6379/0')

# 自定义WebDriver上下文管理器
class WebDriverContext:
    def __init__(self):
        self.driver = None
    
    def __enter__(self):
        self.driver = webdriver.Chrome()
        return self.driver
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        if self.driver:
            self.driver.quit()
            print("WebDriver已关闭")

@app.task
def selenium_task():
    with WebDriverContext() as driver:
        # 业务逻辑
        driver.get("https://example.com")
        # ... 操作

建议把这个方案和前面的信号处理结合使用，作为正常流程的兜底。

方案4：撤销任务时主动查找并关闭Driver #

如果你不想用信号机制，也可以在调用revoke后，主动通过任务ID找到对应的Driver进程并杀死。比如在任务启动时把Driver的进程ID存在Redis，撤销时取出PID并终止进程。

任务代码：

import os
import redis
from selenium import webdriver
from celery import Celery

app = Celery('proj', broker='redis://localhost:6379/0')
r = redis.Redis(host='localhost', port=6379, db=0)

@app.task(bind=True)
def selenium_task(self):
    driver = webdriver.Chrome()
    # 获取Chrome驱动的进程PID
    driver_pid = driver.service.process.pid
    # 存储任务ID与PID的映射
    r.set(f"task_{self.request.id}_driver_pid", driver_pid)
    
    try:
        # 业务逻辑
        driver.get("https://example.com")
        # ...
    finally:
        driver.quit()
        r.delete(f"task_{self.request.id}_driver_pid")

撤销任务时的处理：

from proj.celery import app
import redis
import os
import signal

r = redis.Redis(host='localhost', port=6379, db=0)

def revoke_task_and_close_driver(task_id):
    # 先撤销Celery任务
    app.control.revoke(task_id, terminate=True)
    # 获取Driver的PID
    driver_pid = r.get(f"task_{task_id}_driver_pid")
    if driver_pid:
        driver_pid = int(driver_pid)
        try:
            os.kill(driver_pid, signal.SIGTERM)
            print(f"已杀死WebDriver进程{driver_pid}")
            r.delete(f"task_{task_id}_driver_pid")
        except OSError as e:
            print(f"杀死进程失败: {str(e)}")

总结一下，最推荐方案1+方案3的组合：用Celery信号处理撤销时的清理，用上下文管理器做正常流程的兜底，这样不管任务是正常结束还是被撤销，WebDriver都能被正确关闭。

内容的提问来源于stack exchange，提问作者Muhammad Zubair