Блог

Я пытаюсь создать API, который использует модель Pytorch. Однако, как только я увеличиваю WEB_CONCURRENCY значение выше 1, он создает существенно больше потоков, чем ожидалось, и значительно замедляется, даже при отправке одного запроса.

Пример кода:

api.sh

 export WEB_CONCURRENCY=2

python api.py
  

api.py

 from starlette.applications import Starlette
from starlette.responses import UJSONResponse
from starlette.middleware.gzip import GZipMiddleware
from mymodel import Model


model = Model()
app = Starlette(debug=False)
app.add_middleware(GZipMiddleware, minimum_size=1000)    


@app.route('/process', methods=['GET', 'POST', 'HEAD'])
async def add_styles(request):
    if request.method == 'GET':
        params = request.query_params
    elif request.method == 'POST':
        params = await request.json()
    elif request.method == 'HEAD':
        return UJSONResponse([], headers=response_header)

    print('===Request body===')
    print(params)

    model_output = model(params.get('data', [])) # It is very simplified. Inside there are 
                                                 # many things that are happening, which 
                                                 # involve file reading/writing 
                                                 # and spawning processes with `popen` that 
                                                 # do even more processing. But I don't 
                                                 # think that should be an issue here.

    return model_output


if __name__ == '__main__':
    uvicorn.run('api:app', host='0.0.0.0', port=int(os.environ.get('PORT', 8080)))
  

Когда WEB_CONCURRENCY=1 в api.sh при nvidia-smi запуске отображается только 1 процесс python, а модель использует 1,2 ГБ или видеопамять. Запрос занимает ~ 0.7с

Когда WEB_CONCURRENCY=2 в api.sh в nvidia-smi может быть более 8 процессов python, и они будут использовать более ~ 8 ГБ видеопамяти. Тогда один запрос может занять до 3 секунд, если вам повезет и вы не получите ошибку нехватки памяти.

Я использую Python3.8

Почему Pytorch не использует ожидаемый объем видеопамяти 2,4 ГБ, когда WEB_CONCURRENCY=2 ? И почему это так сильно замедляется?

Если кто-нибудь еще столкнется с этой проблемой, просто используйте gunicorn. Он использует отдельные потоки / процессы, поэтому внутреннего конфликта не происходит.

Поэтому вместо того, чтобы запускать его с помощью: python api.py , просто запустите с: gunicorn -w 2 api:app -k uvicorn.workers.UvicornWorker