将 PyTorch 与 Celery 一起使用答案

【问题标题】：Using PyTorch with Celery将 PyTorch 与 Celery 一起使用
【发布时间】：2020-08-29 09:17:46
【问题描述】：

我正在尝试在 Django 应用程序中运行 PyTorch 模型。由于不建议在视图中执行模型（或任何长时间运行的任务），我决定在 Celery 任务中运行它。我的模型很大，加载大约需要 12 秒，推断大约需要 3 秒。这就是为什么我决定不能在每次请求时都加载它。所以我尝试在设置中加载它并将其保存在那里以供应用程序使用。所以我的最终方案是：

当 Django 应用启动时，会在设置中加载 PyTorch 模型，并且可以从应用访问它。
views.py 收到请求时，会延迟一个 celery 任务
celery 任务使用 settings.model 推断结果

这里的问题是celery任务在尝试使用模型时抛出如下错误

[2020-08-29 09:03:04,015: ERROR/ForkPoolWorker-1] Task app.tasks.task[458934d4-ea03-4bc9-8dcd-77e4c3a9caec] raised unexpected: RuntimeError("Cannot re-initialize CUDA in forked subprocess. To use CUDA with multiprocessing, you must use the 'spawn' start method")
Traceback (most recent call last):
  File "/home/ubuntu/anaconda3/envs/tensor/lib/python3.7/site-packages/celery/app/trace.py", line 412, in trace_task
    R = retval = fun(*args, **kwargs)
  File "/home/ubuntu/anaconda3/envs/tensor/lib/python3.7/site-packages/celery/app/trace.py", line 704, in __protected_call__
    return self.run(*args, **kwargs)
  /*...*/
  File "/home/ubuntu/anaconda3/envs/tensor/lib/python3.7/site-packages/torch/cuda/__init__.py", line 191, in _lazy_init
    "Cannot re-initialize CUDA in forked subprocess. " + msg)
RuntimeError: Cannot re-initialize CUDA in forked subprocess. To use CUDA with multiprocessing, you must use the 'spawn' start method

这是我的 settings.py 中加载模型的代码：

if sys.argv and sys.argv[0].endswith('celery') and 'worker' in sys.argv: #In order to load only for the celery worker
    import torch
    torch.cuda.init()
    torch.backends.cudnn.benchmark = True
    load_model_file()

还有任务代码

@task
def getResult(name):
    print("Executing on GPU:", torch.cuda.is_available())
    if os.path.isfile(name):
        try:
            outpath = model_inference(name)
            os.remove(name)
            return outpath
        except OSError as e:
            print("Error", name, "doesn't exist")
    return ""

任务中的打印显示"Executing on GPU: true"

我尝试在 torch.cuda.init() 之前和之后的 settings.py 中设置 torch.multiprocessing.set_start_method('spawn')，但它给出了相同的错误。

【问题讨论】：

标签： python django multiprocessing pytorch celery

【解决方案1】：

只要您还使用同一库中的Process，设置此方法就可以工作。

from torch.multiprocessing import Pool, Process

Celery 使用“常规”multiprocessing 库，因此出现此错误。

如果我是你，我会尝试：

run it single threaded 看看是否有帮助
run it with eventlet 看看是否有帮助
read this

【讨论】：

不需要“torch.multiprocessing.set_start_method('spawn')”或“from torch.multiprocessing import Pool, Process”。使用 eventlet 和并发 0 解决了问题
单线程也可以工作 - celery -A your_proj worker -P solo 有没有办法知道哪个是更好的解决方案 - 单线程中的单线程与并发 0 的 eventlet/gevent？

【解决方案2】：

一个快速的解决方法是使事情成为单线程的。为此，在启动 celery worker 时将 celery 的 worker pool 类型设置为solo

celery -A your_proj worker -P solo -l info

【讨论】：

【解决方案3】：

这是因为 Celery worker 本身正在使用分叉。这似乎是当前的 known issue 与 Celery >=4.0

您曾经能够将 celery 配置为 spawn，而不是 fork，但该功能 (CELERYD_FORCE_EXECV) 是 removed in 4.0。

没有内置选项可以解决此问题。一些自定义的猴子补丁可能是可能的，但是 YMMV

一些可能可行的选择可能是：

使用 celery <4.0 并启用 CELERYD_FORCE_EXECV。
在 Windows 上启动 celery worker（无论如何都无法进行分叉）

【讨论】：