Celery调用AI服务

Celery调用AI服务

问题引入

​ 在AI模型预测初始化模型网络时,将模型中的参数展到内存中，如果模型过大，每次调用模型时，都对模型参数进行初始化，存在耗时较长的问题，会增加整体预测耗时，如果能设置一个全局变量保存模型参数，保证每次调用时，直接使用模型参数，会不会更好的缩短预测时间

解决方案

​ 如果你在使用Celery调用GPU服务，并且需要将GPU模型的参数保存在全局变量中以便后续任务使用，可以考虑使用Celery的任务上下文(Context)。
Celery的任务上下文是一种机制，用于在任务之间传递数据和状态。通过使用Celery的任务上下文，可以在任务之间共享全局变量，例如GPU模型的参数。对于只需要加载一次并保存在全局变量中的GPU模型，可以在Celery worker启动时加载模型，并将模型参数保存在任务上下文中。后续任务可以从任务上下文中读取模型参数，而无需重新加载模型。
以下是一个示例代码，演示如何在Celery中调用GPU服务，并将GPU模型参数保存在任务上下文中：
pythonCopy code

import torch
from celery import Celery, Task

# 创建Celery应用程序

app = Celery('tasks', broker='pyamqp://guest@localhost//', backend='rpc://')

# 加载模型，并将模型参数保存在任务上下文中

class LoadModelTask(Task):
    def __init__(self):
        super().__init__()
        self.model = None
    

    def run(self):
        self.model = torch.load('/path/to/model.pth')
        self.app.conf.CELERY_CONTEXT['model'] = self.model
        return True

# 在启动Celery worker时注册任务上下文

@app.on_after_configure.connect
def setup_global_context(sender, **kwargs):
    sender.conf.CELERY_CONTEXT = {}

# 定义一个使用GPU模型的任务

@app.task
def use_model(x):
    model = app.conf.CELERY_CONTEXT['model']
    # 在这里使用GPU模型进行计算
    return x

if __name__ == '__main__':

    # 启动Celery worker

​    app.worker_main(argv=['worker', '--loglevel=INFO'])

问题总结

在上述代码中，我们定义了一个名为 LoadModelTask 的Celery任务，用于加载GPU模型并将模型参数保存在任务上下文中。我们还在启动Celery worker时注册了任务上下文。在使用GPU模型的任务 use_model 中，我们从任务上下文中读取模型参数，并在计算过程中使用该模型。
使用这种方法，可以确保GPU模型只被加载一次，并且可以轻松地在任务之间共享模型参数和其他全局变量。