GUnicorn + CUDA: невозможно повторно инициализировать CUDA в разветвленном подпроцессе ⇐ Python

[Anonymous]

Я создаю службу вывода с факелом, пушкой и колбой, которая должна использовать CUDA. Чтобы снизить требования к ресурсам, я использую опцию предварительной загрузки Gunicorn, поэтому модель распределяется между рабочими процессами. Однако это приводит к проблеме с CUDA. Следующий фрагмент кода показывает минимальный пример воспроизведения:

Код: Выделить всё

from flask import Flask, request
import torch

app = Flask('dummy')

model = torch.rand(500)
model = model.to('cuda:0')

@app.route('/', methods=['POST'])
def f():
data = request.get_json()
x = torch.rand((data['number'], 500))
x = x.to('cuda:0')
res = x * model
return {
"result": res.sum().item()
}

Запуск сервера с помощью CUDA_VISIBLE_DEVICES=1 Gunicorn -w 3 -b $HOST_IP:8080 --preload run_server:app позволяет службе успешно запуститься. Однако после выполнения первого запроса (

Код: Выделить всё

curl -X POST -d '{"number": 1}'

), рабочий выдает следующую ошибку:

Код: Выделить всё

[2022-06-28 09:42:00,378] ERROR in app: Exception on / [POST]
Traceback (most recent call last):
File "/home/user/.local/lib/python3.6/site-packages/flask/app.py", line 2447, in wsgi_app
response = self.full_dispatch_request()
File "/home/user/.local/lib/python3.6/site-packages/flask/app.py", line 1952, in full_dispatch_request
rv = self.handle_user_exception(e)
File "/home/user/.local/lib/python3.6/site-packages/flask/app.py", line 1821, in handle_user_exception
reraise(exc_type, exc_value, tb)
File "/home/user/.local/lib/python3.6/site-packages/flask/_compat.py", line 39, in reraise
raise value
File "/home/user/.local/lib/python3.6/site-packages/flask/app.py", line 1950, in full_dispatch_request
rv = self.dispatch_request()
File "/home/user/.local/lib/python3.6/site-packages/flask/app.py", line 1936, in dispatch_request
return self.view_functions[rule.endpoint](**req.view_args)
File "/home/user/project/run_server.py", line 14, in f
x = x.to('cuda:0')
File "/home/user/.local/lib/python3.6/site-packages/torch/cuda/__init__.py", line 195, 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

Я загружаю модель в родительский процесс, и она доступна каждому разветвленному рабочему процессу. Проблема возникает при создании тензора с поддержкой CUDA в рабочем процессе. Это повторно инициализирует контекст CUDA в рабочем процессе, но происходит сбой, поскольку он уже был инициализирован в родительском процессе. Если мы установим x = data['number'] и удалим x = x.to('cuda:0'), вывод будет успешным.
Добавление torch.multiprocessing.set_start_method('spawn') или multiprocessing.set_start_method('spawn') ничего не изменит, возможно, потому, что Gunicorn обязательно будет использовать fork при запуске с - -preload.
Решением может быть отказ от использования параметра --preload, что приводит к созданию нескольких копий модели в памяти/графическом процессоре. Но именно этого я пытаюсь избежать.
Есть ли возможность решить эту проблему без загрузки модели отдельно в каждом рабочем процессе?

Подробнее здесь: https://stackoverflow.com/questions/727 ... subprocess