无法在 Celery 任务中将 Pytorch 与 CUDA 一起使用

Question

我有一个使用 torch 库的 celery 任务，它在内部使用 CUDA。当我运行任务时，它没有说 “无法在分叉子进程中重新初始化 CUDA。要将 CUDA 与多处理一起使用，您必须使用 'spawn' 启动方法”

当我稍微浏览一下时，我得到了这个 - https://github.com/celery/celery/issues/6036 这个issue说Celery只支持fork不支持spawn。

这个有什么workaround/alternative吗？

Answer 1

你应该加载 pytorch 模型并将其用作全局变量。模型和任务的实例将运行在同一个进程

from celery.signals import worker_process_init
# for more information about worker_process_init, read here:
# https://docs.celeryproject.org/en/stable/userguide/signals.html#worker-process-init

pytorch_model = None

@worker_process_init.connect()
def init_worker_process(**kwargs):
    """
    load model before running tasks
    :param kwargs:
    :return:
    """
    global pytorch_model
    pytorch_model = load_model()

@app.task
def predict_task(image: np.ndarray):
    return pytorch_model.predict(image)

无法在 Celery 任务中将 Pytorch 与 CUDA 一起使用

Unable to use Pytorch with CUDA in Celery task

celery

pytorch