FastAPI – Hamel’s Blog - Hamel Husain

FastAPI — это веб-фреймворк для Python. Его любят использовать для обслуживания прототипов ML-моделей.

Впечатления

Фреймворки для обслуживания моделей (TF Serving, TorchServe и т. д.), вероятно, и есть верный путь для продакшена / корпоративных развёртываний, особенно для крупных моделей. Они предлагают больше возможностей, а задержка будет более предсказуемой (пусть и более медленной). Думаю, что для небольших моделей (< 200 МБ) FastAPI вполне подходит. Начать работать с FastAPI очень легко. Мне удалось подтвердить бенчмарк Sayak, где FastAPI быстрее, чем TF Serving, но при этом в целом менее стабилен. FastAPI также с большей вероятностью даёт сбой, хотя мне не удалось его вызвать. В моих экспериментах FastAPI был намного быстрее на этой небольшой модели, но с более крупными моделями ситуация может измениться. Память расходуется линейно по мере увеличения числа воркеров Uvicorn. Фреймворки для обслуживания моделей вроде TF-Serving, похоже, работают эффективнее. Следует не забыть установить переменную окружения TF_FORCE_GPU_ALLOW_GROWTH=true, если вы выполняете инференс на GPU. Думаю, во многих случаях вы будете выполнять инференс на CPU, так что бóльшую часть времени это может быть неактуально. FastAPI выглядит действительно удобным для небольших моделей и ограниченного оборудования, где на узле всего один воркер, а балансировка нагрузки идёт между узлами (потому что вы не реплицируете модель с каждым воркером). Отладка FastAPI — это нечто потрясающее: это чистый Python, и вы получаете удобную страницу документации по адресу http:///docs, которая позволяет прямо на странице протестировать ваши эндпоинты! Документация по FastAPI тоже великолепна. Если вы хотите, чтобы параметры запроса передавались в теле (как это часто бывает в ML, ведь вы хотите отправить данные для оценки), вам придётся использовать Pydantic. Это весьма навязанный подход, но достаточно простой в использовании.

Загрузка модели и получение предсказаний

Будем использовать модель, обученную в руководстве по TF Serving. Более того, мы загрузим её из формата SavedModel.

# this cell is exported to a script from fastapi import FastAPI, status from pydantic import BaseModel from typing import List import tensorflow as tf import numpy as np def load_model(model_path='/home/hamel/hamel/notes/serving/tfserving/model/1'): "Load the SavedModel Object." sm = tf.saved_model.load(model_path) return sm.signatures["serving_default"] # this is the default signature when you save a model

# this cell is exported to a script def pred(model: tf.saved_model, data:np.ndarray, pred_layer_nm='dense_3'): """ Make a prediction from a SavedModel Object. `pred_layer_nm` is the last layer that emits logits. https://www.tensorflow.org/guide/saved_model """ data = tf.convert_to_tensor(data, dtype='int32') preds = model(data) return preds[pred_layer_nm].numpy().tolist()

Тестовые данные

_, (x_val, _) = tf.keras.datasets.imdb.load_data(num_words=20000) x_val = tf.keras.preprocessing.sequence.pad_sequences(x_val, maxlen=200)[:2, :]

Делаем предсказание

model = load_model() pred(model, x_val[:2, :])

[[0.8761785626411438, 0.12382148206233978], [0.0009457750129513443, 0.9990542531013489]]

Создаём приложение FastAPI

# this cell is exported to a script app = FastAPI() items = {} @app.on_event("startup") async def startup_event(): "Load the model on startup https://fastapi.tiangolo.com/advanced/events/" items['model'] = load_model() @app.get("/") def health(status_code=status.HTTP_200_OK): "A health-check endpoint" return 'Ok'

Мы хотим отправлять данные для предсказания в теле запроса (а не через параметры пути). Согласно документации:

FastAPI распознает, что параметры функции, совпадающие с параметрами пути, должны браться из пути, а параметры функции, объявленные как модели Pydantic, должны браться из тела запроса.

# this cell is exported to a script class Sentence(BaseModel): tokens: List[List[int]] @app.post("/predict") def predict(data:Sentence, status_code=status.HTTP_200_OK): preds = pred(items['model'], data.tokens) return preds

Резюме: приложение FastAPI

Давайте посмотрим на main.py со всеми собранными воедино частями:

#This is a hack for Quarto for generated scripts from IPython.display import display, Markdown code = !cat main.py display(Markdown('```{.python filename="main.py"}\n' + '\n'.join(code) + '\n```'))

main.py

# AUTOGENERATED! DO NOT EDIT! File to edit: index.ipynb. # %% auto 0 __all__ = ['app', 'items', 'load_model', 'pred', 'startup_event', 'health', 'Sentence', 'predict'] # %% index.ipynb 3 from fastapi import FastAPI, status from pydantic import BaseModel from typing import List import tensorflow as tf import numpy as np def load_model(model_path='/home/hamel/hamel/notes/serving/tfserving/model/1'): "Load the SavedModel Object." sm = tf.saved_model.load(model_path) return sm.signatures["serving_default"] # this is the default signature when you save a model # %% index.ipynb 4 def pred(model: tf.saved_model, data:np.ndarray, pred_layer_nm='dense_3'): """ Make a prediction from a SavedModel Object. `pred_layer_nm` is the last layer that emits logits. https://www.tensorflow.org/guide/saved_model """ data = tf.convert_to_tensor(data, dtype='int32') preds = model(data) return preds[pred_layer_nm].numpy().tolist() # %% index.ipynb 10 app = FastAPI() items = {} @app.on_event("startup") async def startup_event(): "Load the model on startup https://fastapi.tiangolo.com/advanced/events/" items['model'] = load_model() @app.get("/") def health(status_code=status.HTTP_200_OK): "A health-check endpoint" return 'Ok' # %% index.ipynb 12 class Sentence(BaseModel): tokens: List[List[int]] @app.post("/predict") def predict(data:Sentence, status_code=status.HTTP_200_OK): preds = pred(items['model'], data.tokens) return preds

Запускаем приложение

Мы можем запустить приложение командой:

uvicorn main:app --host 0.0.0.0 --port 5701

main соответствует файлу main.py app соответствует объекту app внутри main.py — app = FastAPI() --reload: заставляет сервер перезапускаться при изменении кода, только для разработки

import requests, json def predict_rest(json_data, url='http://localhost:5701/predict'): json_response = requests.post(url, json={'tokens': json_data}) return json.loads(json_response.text)

predict_rest(x_val.tolist())

[[0.8761785626411438, 0.12382148206233978], [0.0009457750129513443, 0.9990542531013489]]

Нагрузочное тестирование FastAPI

Это действительно быстро

from fastcore.parallel import parallel from functools import partial parallel_pred = partial(parallel, threadpool=True, n_workers=500)

sample_data = [x_val.tolist()] * 1000

%%time results = parallel_pred(predict_rest, sample_data)

CPU times: user 2.29 s, sys: 252 ms, total: 2.54 s Wall time: 2.38 s

Добавляем воркеры Uvicorn

Uvicorn также имеет возможность запускать и работать с несколькими процессами-воркерами. Тем не менее на данный момент возможности Uvicorn по управлению процессами-воркерами более ограничены, чем у Gunicorn. Поэтому, если вы хотите иметь менеджер процессов на этом уровне (на уровне Python), возможно, лучше попробовать использовать Gunicorn в качестве менеджера процессов.

Вы можете добавить воркеры Uvicorn с помощью флага --workers:

uvicorn main:app --host 0.0.0.0 --port 5701 --workers 8

GPU

Когда я масштабировался до 8 воркеров на GPU, я получил ошибки OOM (нехватки памяти). Чтобы этого избежать, нужно ограничить рост использования памяти GPU, установив TF_FORCE_GPU_ALLOW_GROWTH в значение true:

TF_FORCE_GPU_ALLOW_GROWTH=true uvicorn main:app --host 0.0.0.0 --port 5701 --workers 8

Из документации:

По умолчанию TensorFlow занимает почти всю память GPU всех GPU (с учётом CUDA_VISIBLE_DEVICES), видимых процессу.

Это означает, что если вы работаете на GPU и у вас > 1 воркера, вы получите воркеры с ошибками OOM, если не задать эту переменную окружения!

%%time results = parallel_pred(predict_rest, sample_data)

CPU times: user 2.26 s, sys: 294 ms, total: 2.55 s Wall time: 2.34 s

Масштабирование числа воркеров не дало никакого эффекта в этом конкретном случае. Возможно, потому что низкая задержка используемой мной модели недостаточно нагружает пропускную способность.