Внедряем модель c использованием model-serving фреймворков, NVIDIA Triton и Torchserve

Введение

Сегодня мы поговорим о процессе внедрения нейронных сетей и фреймворках сервинга — инструментах, которые значительно его упрощают.

Мы познакомимся с основными концепциями и примерами с кодом, которые помогут глубже погрузиться в тему:

1. Рассмотрим, как и зачем встраивать модель в сервис, а также обсудим важные аспекты проектирования встраиваемых систем и клиент-серверных приложений.
2. Изучим аспекты дизайна клиент-серверного приложения в контексте нейронных сетей и подходы к обработке данных в различных сценариях.
3. Углубимся в принципы работы asyncio и обсудим мотивацию использования сервинг-фреймворков.
4. Рассмотрим общие компоненты и архитектуру сервинг-фреймворков, а также их особенности на примере TorchServe и NVIDIA Triton Inference Server.
5. Разберём применение NVIDIA Triton Inference Server, создадим базовое приложение на его основе и FastAPI, которое можно будет использовать в качестве примера при самостоятельной работе.

Итак, давайте начинать! 🙂

Зачем встраивать модель в сервис?

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

После стадии исследования у нас обычно появляются следующие артефакты:

• код обучения;
• набор залогированных экспериментов;
• чекпоинт модели;
• карточка с описанием метрик, данных и дизайна эксперимента, на основе которого создана модель.

Мы можем провести локальное тестирование модели на этом наборе данных, а также оценить её на внешних тестовых наборах и провести дополнительный анализ её поведения.

Если цель — показать высокие результаты на научном датасете, то на этом этапе наша работа может быть закончена (при удовлетворённости результатом 🙂). Однако в большинстве случаев модель должна быть доступна не только нам, но и другим пользователям.

Поэтому давайте разберём несколько основных направлений использования модели в продуктовой среде.

Использование внутри встраиваемого / программного решения

В этом случае модель будет выступать как функция или часть пайплайна для обработки данных. Например, модель может решать задачу обнаружения внутри некоторого роботизированного комплекса или быть одним из компонентов программного пакета для редактирования изображения. Тогда, как правило, мы сталкиваемся с заранее определённым потоком данных. Значит, мы хотим минимизировать время на обработку данных или, к примеру, сделать модель более отзывчивой к запросу пользователя. Чаще всего в таком случае проектирование сервиса с моделью опирается на следующие подходы:

• использование внутри микросервиса или как часть монолитной архитектуры;
• реализация обёртки вокруг модели на максимально эффективных (C++, Rust) языках для минимизации задержек;
• конвертация модели в один из фреймворков эффективного инференса (OpenVino / TensorRT) или использование проприетарных конвертеров в случае инференса на одноплатных компьютерах;
• минимизация расходов на передачу данных внутри системы.

Использование в рамках клиент-серверного приложения

Если вы ведёте разработку продуктов для множества пользователей, вы сталкиваетесь с необходимостью встраивать модель так, чтобы к ней одновременно (или с небольшим ожиданием) могло обращаться много клиентов. Давайте подробнее рассмотрим построение таких систем.

При дизайне того, как мы встроим модель в подобную систему, мы обращаем внимание на следующие детали:

• Какое количество пользователей у нашей системы?
• Как приходят от них запросы?
• Что важнее — обработать индивидуальные запросы максимально быстро, но поочередно, или сделать так, чтобы все пользователи немного подождали, но при этом получили результат?
• Должны ли различные функции (или модели) в системе работать синхронно, или же нам необходима асинхронная обработка?

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

💡Latency (задержка) — она описывает время с момента отправки запроса пользователем до получения им ответа.

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

💡Throughput (пропускная способность) — количество запросов, обрабатываемых за определённый промежуток времени. Если система обслуживает много пользователей и принимает множество запросов одновременно, то целью является формирование ответов на как можно большее количество запросов за заданное время. В отличие от показателя latency, здесь мы стремимся к максимальным значениям.

Между двумя характеристиками есть зависимость. Давайте рассмотрим следующий пример: предположим, что время, нужное нейронной сети для обработки изображения (latency), составляет 100 мс. Если мы будем обрабатывать каждый пользовательский запрос по отдельности, то сможем обработать 10 запросов в секунду — значение throughput будет равно 10. Задержка от момента передачи данных в сеть до момента выдачи результата также составит 100 мс.

Теперь рассмотрим другой подход к обработке запросов. Операции над тензорами и матрицами, выполняемые на GPU, позволяют значительно увеличивать пропускную способность при обработке данных батчами. Допустим, обработка 10 запросов по одному займёт 1 секунду, а если мы объединим их в батч, то это время сократится до 500 мс. В результате значение throughput увеличится в два раза, и за одну секунду мы сможем обработать не 10, а 20 запросов. Однако задержка (latency) при этом возрастёт в пять раз: от момента формирования батча до получения результата должно пройти 500 мс вместо 100 при обработке запросов по одному.

Таким образом, мы видим — чем больше значение latency, тем выше пропускная способность. Однако следует учесть несколько факторов, которые могут повлиять на эту зависимость:

1. Время сбора данных в батч. Чтобы сформировать батч, требуется некоторое время, в течение которого мы ожидаем поступления данных.
2. Скорость доставки батча до GPU. Обычно мы храним данные для формирования батча в оперативной памяти компьютера. Процесс передачи большого объёма данных из одного процесса в другой может занимать значительное время.
3. Количество данных. Слишком большой батч может не влезть в GPU.

Аспекты проектирования клиент-серверного приложения

Давайте вернёмся к обсуждению деталей, на которые мы обращаем особое внимание при проектировании системы.

Количество пользователей системы

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

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

А если ночью пользователи почти не отправляют запросы, при этом днём их много? В таком случае лучше динамически изменять количество поднятых реплик, чтобы ночью не тратить слишком много ресурсов, а днём не захлебнуться от запросов.

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

Отправление запросов пользователями

Есть несколько распространённых ситуаций, при которых пользователи отправляют данные.

1. Данные поступают с большим фиксированным интервалом

Параметры:

• время обработки запроса — 20 мс;
• количество одновременных запросов — 1 запрос;
• частота поступления запросов — каждые 30 мс.

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

2. Данные поступают одновременно через определённый фиксированный интервал

Параметры:

• время обработки запроса — 20 мс;
• количество одновременных запросов — 5 запросов;
• частота поступления запросов — каждые 30 +-5 мс (+-5 здесь — разброс времени для поступления 5 запросов).

Это нераспространённый сценарий, но он возникает, когда мы имеем дело с потоковыми данными от множества работающих синхронно источников. Чтобы достичь максимальной производительности (throughput), нам нужно объединить данные в пакеты (батчи). Обычно для этого используются механизмы, которые ждут поступления всех данных на сервер в течение определённого интервала и затем отправляют их в обработку в виде одного батча.

3. Данные поступают с фиксированным средним интервалом

Параметры:

• время ожидания обработки данных — 20 мс;
• количество одновременных запросов — 1 запрос;
• частота поступления запросов — каждые 10 мс.

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

В таком случае можно попробовать поднять несколько копий нейронной сети на одном сервере. Копии будут работать конкурентно, позволяя нам распределить запросы пользователей между моделями. Однако у него есть существенный недостаток — каждая копия нейронной сети занимает часть ресурсов GPU. Следовательно, каждая отдельная модель будет работать медленнее. А если все ресурсы GPU будут заняты (например, если их хватает только на три нейронных сети, а мы создали четыре), то какой-то из инстансов будет вынужден ждать освободившихся ресурсов.

Если не удаётся обработать все запросы силами одного сервера — можно распределить запросы между моделями на разных серверах.

4. Данные поступают беспорядочно

Параметры:

• время ожидания обработки данных — 20 мс;
• количество одновременных запросов — 1-20 запросов;
• частота поступления запросов — каждые 20 +-10 мс.

Этот случай — наиболее типичный. В подобных ситуациях рекомендуется применять комбинацию двух подходов:

1. Батчевая обработка: данные, полученные в течение определённого времени, которые группируются в батчи.
2. Распределение запросов: если батчевый обработчик занят, запросы направляются в дополнительные инстансы моделей, которые могут быть подняты автоматически за счёт динамического масштабирования.

Итак, мы рассмотрели несколько типов нагрузки на наше приложение и возможности работы с ними.

Давайте создадим простенькое приложение с нейросетью на FastAPI и обсудим, когда такой подход «пойдёт», а когда стоит обратиться к дополнительным фреймворкам для сервинга нейронных сетей.

Приложение с использованием FastAPI

Рассмотрим простой пример создания приложения на основе FastAPI, куда мы встроим нейронную сеть.

Для наглядности разделим приложение на несколько модулей:

1. Реализация модели с использованием PyTorch.
2. Создание структур данных для запросов и ответов приложения с помощью pydantic.
3. Реализация внутренней логики API на основе FastAPI.
4. Управление серверной частью приложения с помощью unvicorn.

Цель приложения — классификация изображений, отправленных пользователями, на 1000 классов, соответствующих датасету Imagenet.

Опишем структуры данных для обработки запросов:

# data_stuctures.py - Структуры данных нашего приложения
from pydantic import BaseModel, Field
class ImageClassficationRequest(BaseModel):
    """Image Request model."""

    image: str = Field(
        ...,
        title="utf-8 string from a base64 encoded image",
    )
    model_name: str = Field(..., title="Name of model", example="densenet")
    
class PredictionResponse(BaseModel):
    """Prediction Response model."""

    class_name: str = Field(..., title="Class name")
    class_id: str = Field(..., title="Class index")
    logit: float = Field(..., title="logit")

Опишем класс, отвечающий за обработку модели:

# inference.py - код вызова Pytorch модели
import json

import torch
from torchvision import models, transforms

from utils import decode_img

class InferenceModule:

    def __init__(self) -> None:
        self.model =  models.densenet121(pretrained=True).to("cuda:0")
        self.model.eval()
        self.idx2name = json.load(open('imagenet_class_index.json'))
        
    def infer_image(self, img_string: str) -> dict:
			tensor = self.transform_image(img_string)
			with torch.no_grad():
				outputs = self.model.forward(tensor)
	    prob, y_hat = outputs.max(1)
	    predicted_idx = str(y_hat.item())
	    cls_name = self.idx2name[predicted_idx]
	    return {"class_name": cls_name, "class_id": predicted_idx, "logit": prob}
	    
    def transform_image(self, img_string: str) -> torch.Tensor:
	    my_transforms = transforms.Compose(
	       [
	          transforms.Resize(224),
            transforms.ToTensor(),
            transforms.Normalize(
                  [0.485, 0.456, 0.406],
                  [0.229, 0.224, 0.225],
             ),
         ]
      )
      image_pil = decode_img(img_string) # PIL image                                         
	    return my_transforms(image_pil).unsqueeze(0)

Опишем эндпоинт для POST-запросов:

from inference import InferenceModule
from data_structures import PredictionResponse,ImageClassficationRequest,TextClassificationRequest
from fastapi import FastAPI, APIRouter, Depends, Response, HTTPException

app = FastAPI()
inference_engine = InferenceModule()

@app.post("/predict-image", response_model=PredictionResponse)
async def predict_image(data: ImageClassficationRequest) -> PredictionResponse:
    prediction = inference_engine.infer_image(data.image)
    return PredictionResponse.parse_obj(prediction)

Отлично, наше базовое приложение готово!

На этом можно было бы закончить пост, но есть одно «но». Это приложение будет работать только в том случае, если им будет пользоваться только один человек. Асинхронный вызов — часть синтаксиса, а не рабочая процедура. Если к нашему сервису будет обращаться несколько пользователей одновременно, они выстроятся в длинную очередь на ожидание обработки своих данных. Давайте разберёмся, почему так происходит. Для этого поймём, как работает Python и asyncio.

Некоторые аспекты asynаcio

Давайте рассмотрим концепции, которые используются в asyncio.

Прежде всего отметим, что asyncio — инструмент для реализации конкурентной обработки в однопоточном режиме. Давайте разберёмся с этим подробнее.

Есть две основные парадигмы обработки множества задач (запросов пользователей), которые применяются повсеместно, — Parallel execution и Concurrency.

💡Parallel execution (параллелизм) — одновременное исполнение нескольких задач. Для достижения параллелизма необходимо физическое одновременное исполнение задач (multithreading и multiprocessing).

💡Concurrency (конкурентность) — процесс, при котором две и более задачи могут запускаться, выполняться и завершаться в перекрывающиеся периоды времени.

При конкурентном выполнении мы можем запускать несколько задач одновременно, и результаты их выполнения могут приходить в разное время. Но если нам нужно выполнить задачу, то выполняться будет только одна она, остальные остановятся. Есть также задачи, выполнение которых позволяет переключаться на другие без остановки (об этом поговорим далее).

С точки зрения реализации кода приложения, представленного в предыдущем разделе, POST-запрос в виде функции def predict_image определён с помощью async def. Так мы определили не просто функцию, а корутину.

💡Корутина (coroutine) — подпрограмма (функция), которую можно приостанавливать и возобновлять без потери состояния.

Все вызовы функций в нашем приложении, определённые таким образом, будут не просто выполняться, а попадать в event_loop — цикл событий, который управляет выполнением всех задач. Эта конструкция не только запускает сами задачи, но и в местах кода с await переключается на другую задачу.

Асинхронные функции подразумевают следующее: каждая задача ожидает время перед выполнением. Например, мы ждём поход во внешнее API и не тратим ресурсы процессора. При этом само время выполнения задачи (если из него вычесть ожидание) довольно маленькое, и на него не тратится много времени процессора. Тогда между задачами можно эффективно переключаться. Пока задача ждёт — она отдаёт управление, чтобы можно было выполнить другие задачи.

На практике применение асинхронных функций особенно оправданно в клиент-серверных приложениях, так как мы часто сталкиваемся со следующими задачами:

1. Авторизация пользователей — обработка ввода-вывода и передача пользовательских данных по сети.
2. Поход в стороннее API.
3. Запись в файлы.

Можно привести много примеров подобных операций, но главное — все они соответствуют нашим критериям:

1. Передача данных, запросы к API или запись в файл требуют определённого времени для выполнения.
2. Они выполняются не самой программой на Python и не «кушают» процессор.
3. Пока мы ожидаем их завершения, мы можем свободно переключаться и запускать другие операции.

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

Все эти задачи можно отнести к категории I/O-Bound. Но если вернуться к задачам, которые решает наше приложение с моделью, все они окажутся CPU- или GPU-Bound. При вызове корутины с таким кодом до завершения её выполнения все остальные корутины будут ожидать своей очереди в event loop. Следовательно, приложение не сможет обрабатывать другие запросы до тех пор, пока корутина не завершит свою работу.

Поэтому когда мы интегрируем нейронные сети в такие приложения, мы используем фреймворки для сервинга нейронных сетей. Они помогают нам развернуть нейронную сеть в рамках отдельного приложения, так что задача корутины по обработке данных пользователя сводится лишь к отправке и приёму данных, а вся сложная обработка выполняется на стороне serving-сервера.

Ключевые компоненты serving-фреймворков

Serving-фреймворки можно описать с помощью следующей схемы:

Фреймворки для сервинга обычно состоят из трёх компонентов:

1. Model serving runtime — важная часть для запуска уже обученных нейронных сетей. В основе могут использоваться различные инструменты и функции для выполнения модели в одном или нескольких фреймворках, например, OnnxRuntime, Torch или TensorRT. Она также отвечает за загрузку модели из хранилища артефактов, где они хранятся. В сервинг-фреймворках процесс конфигурации осуществляется через конфигурационные файлы, в которых описываются параметры входных и выходных данных, а также указываются различные настройки (и оптимизации, применяемые к модели).
2. Model Server
   • Model Server в таких фреймворках обычно отвечает за приём данных и передачу их в нужную модель. Он также собирает данные в пакеты, создаёт и удаляет дополнительные копии модели.
   • API Gateway — часть архитектуры фреймворков, которая обычно называется «Client», или «Frontend». Через API Gateway можно отправлять данные и получать результаты работы моделей с помощью протоколов HTTP или gRPC. Также можно получать статистику работы и влиять на работу сервера через сервисные сообщения.
   • Load Balancer — уже знакомый балансировщик нагрузки. У каждого фреймворка есть свои функции, но концепция остаётся неизменной.
3. Monitoring Metrics — компонент фреймворка для сбора метрик. Обычно в стандартной конфигурации логгируются такие метрики, как Latency, Throughput, а также Uptime моделей. На практике список метрик можно расширять. Как правило, предполагается, что за обработку метрик и визуализацию будут отвечать внешние решения, например, связка Prometheus + Grafana.

TorchServe и NVIDIA Triton Inference Server как serving-фреймворки

Рассмотрим два популярных фреймворка для сервинга моделей: TorchServe и Triton Inference Server.

В качестве стартового решения будет рассмотрен TorchServe, а в качестве наиболее полного решения — NVIDIA Triton Inference Server. Отметим их основные особенности, плюсы и минусы использования.

TorchServe

TorchServe — часть экосистемы Pytorch, разработанная для обеспечения плавного перехода от обучения моделей к их инференсу. На верхнем уровне его работа соответствует общему описанию модели из предыдущего раздела.

Особенности:

1. Фреймворк построен на двух языках: часть, отвечающая за взаимодействие с внешними запросами, написана на Java, а серверная часть — на Python и PyTorch (ещё есть экспериментальный бэкенд на C++).
2. Для загрузки моделей фреймворк использует формат mar (идентичен zip-архиву), который создаётся с помощью утилиты torch-model-archiver. Модель состоит из нескольких компонентов, ниже перечислены основные:
   • handler — класс на Python, который декодирует входные данные и кодирует выходные для отправки в качестве ответа. Он может быть как дефолтным, так и кастомизированным под свои нужды;
   • checkpoint модели в формате torch или torchscript (также поддерживается ONNX);
   • список библиотек или дополнительных файлов (при необходимости);
3. Можно сконфигурировать настройки батчевой обработки и количество воркеров с моделью.
4. Можно настроить последовательный или параллельный порядок выполнения скриптов и моделей.

Преимущества:

1. Простота работы с Python-кодом и моделями на основе Pytorch. Чтобы создать готовую модель, достаточно «заполнить» несколько методов у базового класса.
2. Низкий порог входа. Освоение фреймворка не требует глубоких знаний и опыта.
3. Удобный подход к созданию многоступенчатых конвейеров с различными моделями.

Недостатки:

1. Отсутствие детальных технических описаний в документации. Для понимания внутренних механизмов необходимо обратиться к исходному коду.
2. Сложность отладки Java-части фреймворка. Тем, кто не знаком с Java, здесь будет непросто.
3. Ограничения скорости обработки данных при конкурентном использовании. Поскольку серверная часть написана на Python, одновременное использование GPU несколькими моделями невозможно. Для решения этой проблемы авторы разработали C++-версию сервера, однако для создания собственного hadler также потребуется написание кода на C++ и предварительное преобразование модели с помощью TorchScript.

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

NVIDIA Triton Inference Server

Фреймворк предлагает широкие возможности для развёртывания различных моделей независимо от того, на каком языке они написаны. В настоящее время его можно считать лучшим выбором, если стоит задача запустить как можно больше типов моделей в рамках одной системы обслуживания моделей.

Особенности:

1. Поддержка широкого набора бэкендов в рамках Model serving runtime:
   • базовые фреймворки — Pytorch / Tensorflow / Onnx;
   • использование бэкендов для дополнительной оптимизации — OnnxRuntime, OpenVino, TensorRT;
   • возможность инференса LLM за счёт интеграции с VLLM и TensorRT-LLM;
   • запуск классических алгоритмов ML и инференса бустингов;
   • возможность кастомизировать логику работы под нестандартные сценарии;
2. Эффективная реализация кода — поддержка как батчевого, так и конкурентного режима исполнения моделей на GPU, надстройки над протоколами передачи данных для использования не только сетевого подключения, но и shared memory.
3. Возможность кастомизировать логику работы для нестандартных сценариев:
   • работа с моделями со скрытым состоянием;
   • батчевая обработка данных с разными размерами в пакетах;
   • обработка очереди сообщений, составляющих последовательность в рамках одной выделенной модели.
4. Возможность реализации бизнес-логики в виде пайплайнов посредством комбинации обращений к серверу и промежуточного python-кода.
5. Клиентская часть в виде библиотеки на языке Python с большим количеством примеров под разные кейсы применения.

Преимущества:

1. Поддержка широкого числа бэкендов. Является лучшим выбором для сервинга моделей в различных фреймворках.
2. Производительность. Демонстрирует высокую производительность, особенно при использовании графических процессоров (GPU).
3. Хорошая документация. Предлагает обширную документацию и множество примеров для практического применения.

Недостатки:

1. Необходимость погружения. Репозиторий фреймворка включает 36 отдельных репозиториев, что требует глубокого понимания для полного ознакомления.
2. Недоработки по LLM. Функционал для сервинга LLM-моделей не всегда хорошо документирован.

Итак, познакомившись с двумя наиболее популярными фреймворками для сервинга, давайте кратко рассмотрим устройство Triton Inference Server’а.

Ключевые моменты использования NVIDIA Triton Inference Server

Запуск

Фреймворк поставляется в виде готового docker-образа, где предустановлен весь необходимый функционал. Также можно дополнительно установить нужные python-библиотеки, если модели представлены в виде python-кода.

CLI

docker run --rm --net=host -v ${PWD}/model_repository:/models nvcr.io/nvidia/tritonserver:24.12-py3 tritonserver --model-repository=/models\\

Docker-compose

version: "3.2"
services:
  image: nvcr.io/nvidia/tritonserver:24.12-py3
  command: tritonserver --model-repository=/models
  ipc: "host"
  pid: "host"
  ports:
     - "8000:8000"
     - "8001:8001"
     - "8002:8002"
  shm_size: '1024mb'
  volumes:
     - ${PWD}/models_repository:/models

Подгрузка моделей

Для запуска модели на сервере с помощью NVIDIA Triton Inference Server нам нужны:

• чекпоинт модели, созданной на одном из поддерживаемых фреймворков;
• информация о входном слое модели, его названии и размерностях, а также о выходных слоях;
• конфиг модели;
• репозиторий, организованный по определённой структуре.

Первые два пункта вполне понятны: нам нужно обучить модель и, например, визуализировать её с помощью netron.ai, чтобы узнать названия и размерности необходимых нод. А вот с третьим и четвёртым пунктами стоит разобраться подробнее. Для начала рассмотрим структуру хранилища моделей:

├─── ...
└───<model_repo>
    └───<model_name>
        ├───config.pbtxt
        └───<version>
            └───model.py

• <model_repo> — папка, в которой хранятся все наши модели;
• <model_name>config.pbtxt — конфигурационный файл, который описывает детали запуска конкретной модели (мы подробнее рассмотрим его позже);
• <version> — папка для нашей модели, название которой должно быть целым числом, начиная с 1;
• model.py — код нашей модели, если в качестве бэкенда используется Python. Если выбран другой фреймворк — вместо него может быть чекпоинт модели, например, model.pt, model.onnx или model.engine.

Описание конфигурационного файла модели

Для описания запуска модели возьмём в качестве примера модель densenet в фреймворке ONNX:

name: "densenet_onnx"
platform: "onnxruntime_onnx"
max_batch_size: 8 
dynamic_batching { 
    ### preferred_batch_size: 4 # только для TensorRT
    max_queue_delay_microseconds: 200000
}
instance_group [ 
  {
    count: 3
    kind: KIND_GPU
    gpus: [0]
  }
input [
  {
    name: "data_0"
    data_type: TYPE_FP32
    format: FORMAT_NCHW
    dims: [ -1,3, 224, 224 ]}
  }
]
output [
  {
    name: "fc6_1"
    data_type: TYPE_FP32
    dims: [ -1,1000 ]
  }
]

Давайте рассмотрим все ключевые настройки:

1. name: название модели, которое должно совпадать с названием её содержащей папки.
2. platform: в NVIDIA Triton Inference Server платформа обозначает бэкэнд, который будет использоваться для вычислений графа модели. В данном случае это OnnxRuntime.
3. max_batch_size: параметр максимального размера батча, с которым будет работать модель. Если при создании модели в фреймворке инференса для 0-й размерности было установлено значение -1, то его возможно изменить на любой желаемый размер (однако при использовании TensorRT есть ограничение, о котором мы поговорим позже).
4. dynamic_batching: настройка для определения работы модуля накопления данных.
   1. max_queue_delay_microseconds: данные накапливаются в течение 200 000 микросекунд с момента получения первого изображения. Если до этого накопится пакет из четырёх изображений — данные отправятся в модель.
   2. preferred_batch_size: если мы используем модель во фреймворке TensorRT, мы выбираем между минимальным и максимальным размером пакета. За пределы этого значения мы выйти не можем — это ограничение, установленное max_batch_size при создании графа. Однако мы можем выбрать желаемый размер из диапазона от 1 до max_batch_size и указать его как preferred_batch_size.
5. instance_group — количество копий (инстансов) модели. В этом случае мы развернём три модели на GPU с индексом 0.
6. input и output — словарь входных и выходных нод модели, к которым фреймворк будет обращаться для подачи данных и получения результатов.
   1. name — имя ноды.
   2. data_type — тип данных. В Triton Inference Server доступны различные типы данных, включая int8 и bf16 (однако не все бэкенды их поддерживают).
   3. dims — размерности тензоров. Если мы указываем значение -1 — значение может быть любым.

Запуск клиентской части фреймворка

Как было сказано ранее, фреймворк позволяет использовать два протокола передачи данных — HTTP и GRPC. Возьмем для примера GRPC, чтобы разобраться подробнее.

Для создания объекта класса клиента нам необходимо установить библиотеку tritonclient, а после ипортировать класс grpcclient. Далее мы можем создать объект данного класса, указав в конструкторе адрес: порт для общения с сервером.

import tritonclient.grpc.aio as grpcclient
url = os.environ.get("TRITON_SERVER_URL", "localhost:8001")
triton_client = grpcclient.InferenceServerClient(url=url)os.environ.get("TRITON_SERVER_URL", "triton:8001")

Подготовка данных к передаче

Для формирования входных и приёма выходных сообщений здесь реализованы такие классы, как InferInput и InferRequestedOutput.

Для создания объектов классов нам нужно указать название входной и выходной ноды (исходя из конфигурации модели), а также размер и тип данных. Ниже приведён пример для модели densenet:

input_node_name = "input:0"
output_node_name = "output:0"
shape = [1,3,224,224]
dtype = TYPE_FP32

inputs = [grpcclient.InferInput(input_node_name, shape, dtype)]
outputs = [grpcclient.InferRequestedOutput(output_node_name)]

После создания объектов входа и выхода необходимо заполнить обьект Inputs данными:

 inputs[0].set_data_from_numpy(img.astype(np.float32))

InferInput принимает данные в виде numpy-массива, однако его элементами могут быть не только числа, но и строки, а также байтовое представление данных.

Отправка данных

Для отправки данных на сервер есть несколько вариантов:

• синхронный — ждём выдачу результата, находясь в текущем контексте;
• асинхронный без ожидания — можем отправить данные, выполнить произвольный код и потом получить результат, но проверка наличия результата ложится на плечи пользователя (правильнее сказать — на программиста 🙂);
• асинхронный с использованием asyncio — наиболее подходящий вариант под клиент-серверное приложение.

Отправка данных:

model_name = "densenet_onnx"
results = await triton_client.infer(
            model_name=model_name,
            inputs=inputs,
            outputs=outputs
        )

Получение результата

outputs = results.as_numpy(output_node_name)

Итак, в этом разделе мы рассмотрели базовое использование NVIDIA Triton Inference Server. Давайте применим его на практике в рамках улучшения клиент-серверного приложения с FastAPI.

Клиент-серверное приложение с FastAPI и Triton Inference Server

Описание конфигурационного файла

name: "densenet_onnx"
platform: "onnxruntime_onnx"
max_batch_size: 8 
dynamic_batching { 
    ### preferred_batch_size: 4 # только для TensorRT
    max_queue_delay_microseconds: 200000
}
instance_group [ 
  {
    count: 3
    kind: KIND_GPU
    gpus: [0]
  }
input [
  {
    name: "data_0"
    data_type: TYPE_FP32
    format: FORMAT_NCHW
    dims: [ -1,3, 224, 224 ]
   }
]
output [
  {
    name: "fc6_1"
    data_type: TYPE_FP32
    dims: [ -1,1000 ]
    label_filename: "densenet_labels.txt"
  }
]

Описание класса, отвечающего за обработку модели

# inference.py - код общения с Triton Inference Server
import os

import grpc
import numpy as np
import tritonclient.grpc.aio as grpcclient

from tritonclient.grpc import service_pb2, service_pb2_grpc
from tritonclient.utils import triton_to_np_dtype

from utils import decode_img

class InferenceModule:

    def __init__(self) -> None:
        """Initialize."""
        self.url = os.environ.get("TRITON_SERVER_URL", "triton:8001")
        self.triton_client = grpcclient.InferenceServerClient(url=self.url)

    async def infer_image(
        self,
        img: str,
        model_name: str = "densenet_onnx",
        top_k=1,
    ) -> dict:

        model_meta, model_config = self.parse_model_metadata(model_name)
        shape = model_meta.inputs[0].shape
        channels, width, height = shape[1:]
        dtype = model_meta.inputs[0].datatype

        img = self.preprocess_image(img, height, width, dtype)

        inputs = [grpcclient.InferInput(model_meta.inputs[0].name, [channels, width, height], dtype)]
        inputs[0].set_data_from_numpy(img.astype(np.float32))

        outputs = [grpcclient.InferRequestedOutput(model_meta.outputs[0].name, class_count=top_k)]

        results = await self.triton_client.infer(
            model_name=model_name,
            inputs=inputs,
            outputs=outputs,
        )
        output = results.as_numpy(model_meta.outputs[0].name)[0]
        cls = "".join(chr(x) for x in output).split(":")
       
        return {"class_name": cls[2], "class_id": cls[1], "logit": cls[0]}
       
    def parse_model_metadata(self, model_name: str) -> object:
        channel = grpc.insecure_channel(self.url)
        grpc_stub = service_pb2_grpc.GRPCInferenceServiceStub(channel)
        metadata_request = service_pb2.ModelMetadataRequest(
            name=model_name
        )
        metadata_response = grpc_stub.ModelMetadata(metadata_request)

        config_request = service_pb2.ModelConfigRequest(
            name=model_name
        )
        config_response = grpc_stub.ModelConfig(config_request)

        return metadata_response, config_response
        
    def preprocess_image(self,
                         img: str,
                         height: int,
                         width: int,
                         dtype="FP32",
                         norm_type="IMAGENET") -> np.ndarray:
        img = decode_img(img) # PIL image
        resized_img = img.resize((width, height))

        np_img = np.array(resized_img)
        npdtype = triton_to_np_dtype(dtype)
        typed = np_img.astype(npdtype)

        if norm_type == "INCEPTION":
            scaled = (typed / 127.5) - 1
        elif norm_type == "IMAGENET":
            scaled = typed - np.asarray((123, 117, 104), dtype=npdtype)
        else:
            scaled = typed

        ordered = np.transpose(scaled, (2, 0, 1))
        return ordere

Описание эндпоинта для POST-запросов

Будем использовать асинхронный вызов, чтобы обрабатывать больше одного пользовательского запроса одновременно:

from inference import InferenceModule
from data_structures import PredictionResponse,ImageClassficationRequest,TextClassificationRequest
from fastapi import FastAPI

app = FastAPI()
inference_engine = InferenceModule()

@app.post("/predict-image", response_model=PredictionResponse)
async def predict_image(data: ImageClassficationRequest) -> PredictionResponse:
    prediction = await inference_engine.infer_image(data.image,model_name="densenet_onnx",top_k=1)
    return PredictionResponse.parse_obj(prediction)

Ура, приложение готово! После интеграции Triton Inference Server’а мы получили приложение, которое может работать со множеством пользователей и запросов от них!

Заключение

В этом материале мы рассмотрели, как интегрировать модель с помощью сервинг-фреймворков. Эта задача имеет множество аспектов: как архитектурных, часть из которых мы обсудили, так и практических. Мы привели примеры реализации как на основе FastAPI и Pytorch, так и с использованием NVIDIA Triton Inference Server.

Важно отметить, что мы рассмотрели только верхушку айсберга, не углубляясь в детали оптимизации полученного пайплайна и использования различных низкоуровневых функций NVIDIA Triton Inference Server. Надеемся, что материал поможет вам сделать первый шаг в решении задачи интеграции модели, а также лучше понять особенности её инференса в рамках клиент-серверных приложений!