Любой МЛщик когда-то задается вопросом: «А как мне деплоить свои модели?». Мы давно нашли на него ответ: это Triton Inference Server, разработка Nvidia. У него полно разных достоинств:

ядро написано на плюсах, а пользоваться можно на питоне;
всеядный — переварит модели на Torch, ONNX, OpenVINO, TensorFlow, Scikit-learn, vLLM и просто питоновские скрипты.
поддерживает инференс нескольких моделей одновременно;
может объединять модели в последовательности (ансамбли);
API-интерфейс по HTTP и gRPC;
всякие фичи для ускорения инференса, типа динамического батчинга.

Далее мы расскажем необходимый минимум о Тритоне и как мы его используем в нашем проекте, детально вы прочитать про него в официальном руководстве на Гитхабе, которое вас проведет через полный цикл создания сервиса. Ссылка на репу с нашим проектом: tritoned_bert.

Минимальная конфигурация сервиса

Разработчики Тритона явно старались сделать так, чтобы со стороны пользователей нужно было минимум усилий, чтобы модель подружилась с сервисом. Разберем тот минимум шагов, что вам нужно сделать.

В корневой папке проекта создадим папку и назовем ее model_repository. В ней создадим еще одну папку с названием модели. Это же название будет в API. Внутри папки модели создадим еще папку с названием 1 и файл config.pbtxt. В папке 1 будет храниться сама модель, а цифра в названии указывает на ее версию. Минимальное содержание конфига выглядит так:

platform: "tensorrt_plan"
max_batch_size: 8
input [
  {
    name: "input0"
    data_type: TYPE_FP32
    dims: [ 16 ]
  },
]
output [
  {
    name: "output0"
    data_type: TYPE_FP32
    dims: [ 16 ]
  }
]

То есть вам необходимо указать:

фреймворк(платформу) модели;
максимальный размер батча — от этой настройки зависят другие механизмы Тритона и это не совсем то же самое, что размер батча при обучении;
конфигурации входных и выходных данных, которые представляет собой списки словарей.

Как не трудно догадаться, ключи в конфигурации данных:

name — имя конкретного входа/выхода. Именно оно будет фигурировать в API.
data_type — тип данных, куда без него;
dims — размерность входного/выходного вектора (тензора, в общем виде)

Кладете модель в папку model_repository/your_model_name/1. Всё, можете запускать:

$ docker run --gpus=all -it --shm-size=256m --rm -p8000:8000 -p8001:8001 -p8002:8002 -v $(pwd)/model_repository:/models nvcr.io/nvidia/tritonserver:22.12-py3 

Чтобы правильно выбрать базовый докер-образ для вашей GPU, если вы хотите ее использовать, смотрите вот это руководство.

Запуск питоновских скриптов на стороне Тритона

Модель-то мы развернули, только вот на вход она принимает тензоры и на выход отдает тоже тензоры. Если мы деплоим какой-нибудь Берт для классификации текстов, то нам сначала надо токенизировать текст, чтобы превратить его в векторы, а над логитами модели мы должны сделать хотя бы argmax, чтобы узнать итоговый класс. Чтобы не заставлять каждого клиента сервиса тащить на себе эту логику, можно сделать питоновскую «модель», которая это будет делать на стороне Тритона.

Чтобы провернуть такой фокус, вам нужно определить вот такой шаблонный класс

class TritonPythonModel:
    def initialize(self, args):
        ...
    def execute(self, requests):
        ...
    def finalize(self):
        ...

Названия методов говорят сами за себя. В деле реализации модели вам очень поможет пакет pb_utils, в котором есть разные функции по преобразованию типов — чувствуется влияние плюсов — и созданию специальных тритоновских объектов, например, тензоров. Кроме того, для метода execute есть разные условия, типа сколько реквестов пришло, столько быть отдано респонсов. Подробнее читайте об этом в документации.

Вот пример реализации execute для нашей модели постпроцессинга предсказаний, где выполняется argmax и происходит маппинг индекса класса в его название.

def execute(self, requests):
    responses = []

    for request in requests:
        in_0 = pb_utils.get_input_tensor_by_name(
            request, "logits"
        ).as_numpy()

        predicts = in_0.argmax(axis=1)
        predicts = np.array([self.id2label[x] for x in predicts], dtype=object)
        out_tensor_0 = pb_utils.Tensor("predicts", predicts)

        inference_response = pb_utils.InferenceResponse(
            output_tensors=[out_tensor_0]
        )
        responses.append(inference_response)
    return responses

Надо подчеркнуть, что мы пока сделали самую базовую реализацию, здесь даже нет какой-то обработки ошибок и логирования.

Ансамбли

Хорошо, есть у нас две «модели» на пре- и постпроцессинг и сама Бертовая модель. Как нам это заставить все работать вместе? Для этого в Тритоне есть ансамбли — эдакие метамодели, которые могут запускать обычные модели в определенном порядке.

Структура у них такая же: в папке с названием ансамбля должен быть файл config.pbtxt и пустая папка 1. Отличается содержание конфига. В нем, кроме описания входов и выходов, надо описать последовательность запуска моделей ensemble_scheduling. Обязательно также указывать маппинг названий входов и выходов, даже если они совпадают.

Как это выглядит такой конфиг для нашего сервиса, можете посмотреть по ссылке — он достаточно большой, чтобы лепить его сюда.

Наш шаблон тритоновских сервисов

Мы хотим, чтобы в нашу систему было супер-просто интегрировать сторонние модели. Чтобы это сделать, мы разработали шаблон, который позволит вам превратить вашу бертовую модель для классификации в тритоновский сервис. Всё, что вам нужно, это сохраненная модель в формате Hugging Face, наш репозиторий и пара зависимостей. Репозиторий мы открыли, можете пробовать. Помните, что он еще в разработке, поэтому могут быть оказии. Если что, создавайте карточку или пишите напрямую @Astromis в тг.

Пользоваться шаблоном просто: вам нужно запустить скрипт ./make_triton_image.sh, передав ему следующие параметры:

путь до модели в формате Hugging Face;
путь до токенизатора;
путь до словаря, в котором ключами являются индексы классов, а значениями — названия классов.

Кроме того, есть опциональные параметры:

model_name — имя модели. Если не передать, то будет ensemble model. Также будет называть контейнер.
container_tag — на самом деде, конечно, тег докер-образа. По умолчанию, latest.
max_batch_size — максимальный размет батча, который будет обрабатывать сервер.

Далее с помощью магии bash, шаблоны превратятся в конкретные настройки, ваша модель преобразуется в ONNX и всё это зашьётся в докер-образ Тритона с названием tritoned_[model_name]:[container_tag]. Важно сказать, что пока базовый докер-образ не изменяется, потому что шаблон построен так, что сервис будет работать на CPU. Мы добавим возможность выбирать режим, как и базовый контейнер, позже.

Запускать контейнер можно так:

$ docker run --shm-size=256m -p8000:8000 -p8001:8001 -p8002:8002  tritoned_[model_name]:[container_tag]

Протестить можно так:

$ curl -X POST http://127.0.0.1:8000/v2/models/ensemble_model/infer -d '{"inputs":[{"name":"text_input","shape":[1,1],"datatype":"BYTES","data":["тестовый тест"]}]}'

Кто такой этот ваш ONNX?

Что это вообще такое? Это Open Neural Network Exchange — открытый стандарт, для представления архитектуры нейросетей. Его поддерживают, наверное, все фреймворки для нейросеток, которые вам могут прийти в голову. «Поддерживают» значит, что фреймворки могут экспортировать объект сетки в этот формат. С тем же успехом, ее можно импортировать в эти фреймворки. Если представить ситуацию, что вам в руки попала какая-то onnx-модель, то вы сможете посмотреть ее структуру с помощью netron, закодить ее на любом фреймворке, поддерживающем ONNX, и использовать ее.

Кроме переносимости у ONNX есть еще одна фишка — собственная среда выполнения onnxruntime, которая позволяет даже на обычном CPU инференсить модели с большей скоростью, в сравнении с исходными фреймворками. А ведь можно еще задействовать аппаратные возможности. На Хабре есть пост, где можно наглядно посмотреть сравнение скоростей. В добавок, инструменты ONNX позволяют просто и быстро квантизировать модель, если готовы пожертвовать качеством ради скорости.

Может быть, вы думаете, что экспорт в ONNX сложный? Буквально одна строчка, что в tansfromers (точнее optimum):

ort_model = ORTModelForSequenceClassification.from_pretrained(model_path, export=True) 

что в torch:

torch.onnx.export(model, ...)

Эту строчку мы и используем для конвертации в нашей утилите utils/convert_to_onnx.py.

Кстати говоря, утилита поймет не только локальный путь, но адрес в Hugging Face Hub и Clearml. Если хотите, то легко сможете добавить другое хранилище, например, WandB, Artifactory, MLFlow и т.д. Модель и токенизатор нужно прописывать отдельно. С одной стороны, неудобно, но зато токенизатор и модель могут лежать в разных местах. Это актуально, когда вы много файнтюните одну базовую модель. В таком случае каждый раз сохранять токенизатор избыточно, достаточно просто указывать токенизатор базовой модели. Пример команды для конвертации:

$ python utils/convert_to_onnx.py -m models/best_model -t "deeppavlov/RuBert" --where_model local --where_tokenizer hf

На этом всё, на связи.