redis celery 工装卸模型

vshtjzan  于 12个月前  发布在  Redis
关注(0)|答案(1)|浏览(169)

我目前有一个使用Celery在耳语AI模型上启动任务的系统。然而,现有的设置需要在每个任务中加载模型,这是次优的,因为重复加载过程消耗了大量的时间。不幸的是,我不能保持模型连续加载,因为其他系统需要访问GPU VRAM。
为了解决这个问题,我正在考虑实现一种系统,在接收到任务时加载模型,在没有任务时卸载模型。这种方法旨在通过最小化模型在内存中加载的时间来优化资源利用率,确保GPU资源的有效使用。我相信这种调整将导致整体系统性能和响应能力的提高。
这是一个想法(显然不起作用)

from celery import Celery, Task
from celery.signals import (
    task_received,
    celeryd_after_setup,
    task_success,
    task_failure,
)
import redis
import torch
import whisper
r = redis.Redis(host=hostname)

model = None

def checkActivesTasksWorker():
    global model
    if r.llen(activesTasksWorker) == 0:
        # Model is deleted when there is no task left on worker
        del model
        torch.cuda.empty_cache()

@celery.task
def myTask():
    launchAnalyticsWhisper(model)

@task_received.connect
def taskReceivedHandler(sender, request, **kwargs):
    if r.llen(activesTasksWorker) == 0:
        model = whisper.load_model("medium")
    r.lpush(activesTasksWorker, request.id)

@task_success.connect(sender=myTask)
def taskSuccessHandler(sender, result, **kwargs):
    r.lrem(activesTasksWorker, 1, result["taskId"])
    checkActivesTasksWorker()

@task_failure.connect(sender=myTask)
def taskFailureHandler(sender, task_id, exception, **kwargs):
    r.lrem(activesTasksWorker, 1, task_id)
    checkActivesTasksWorker()

@celeryd_after_setup.connect
def initList(sender, instance, **kwargs):
    # Clear the active tasks for the worker
    r.delete(activesTasksWorker)

字符串
model处处都是None,所以全局变量不起作用。你有办法让它起作用吗?

wb1gzix0

wb1gzix01#

Celery为每个任务生成新的进程,这就是为什么它们不能访问全局变量。
最好的方法是为所讨论的模型提供一个专用的推理服务器。推理服务器在启动时加载模型。celery任务向推理服务器发出请求。
如果您必须在同一个GPU上处理多个模型,您可以向推理服务器添加一个端点,以便在CPU内存和GPU内存之间移动模型。这样,您可以保持GPU内存空闲,同时仍然只从磁盘加载模型一次。

相关问题