Магнит рассказал, как задеплоил векторный поиск на Triton без GPU

Команда поиска MAGNIT OMNI перевела сервис векторного поиска на NVIDIA Triton Inference Server и отказалась от инференса на Python. Об этом в блоге на Хабре рассказал ведущий ML-разработчик Стас Чернышев. По его словам, инференс одного запроса на CPU через FastAPI занимал около 30 мс, что не устраивало команду.

Почему отказались от FastAPI

FastAPI — привычное решение для обёртки моделей в REST, но у него нет инфраструктурных функций из коробки: gRPC, батчинга входных запросов, кэша, health check и метрик. Всё это приходится писать и поддерживать вручную. Кроме того, переход с CPU на GPU требует отдельной работы с драйверами в Dockerfile.

Разработчики сравнили FastAPI, BentoML и Triton по набору возможностей. Динамический батчинг есть у BentoML и Triton, кэш ответов — только у Triton (локальный или на Redis), метрики Prometheus — у BentoML и Triton. Инференс без Python поддерживает только Triton за счёт нативных бэкендов. В BentoML инференс идёт в Python-процессах, а сами разработчики для ускорения предлагают запускать внутри него Triton.

Как устроена модель

Базовая модель эмбеддингов — BERTA, компактная русскоязычная модель, полученная дистилляцией эмбеддингов FRIDA от SaluteDevices. Она наследует систему префиксов FRIDA, в том числе «search_query:» и «search_document:» для асимметричного поиска. По этому признаку команда разделила функционал на два сервиса: эмбеддер запросов и эмбеддер документов.

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

Отказ от Python

На Python осталась только токенизация: небольшая модель на Python backend добавляет префикс, вызывает токенизатор и передаёт тензоры дальше. Тело модели сконвертировали в ONNX и отдали нативному onnxruntime backend. В репозиторий моделей достаточно положить веса model.onnx и config.pbtxt — остальное Triton делает сам.

Модель обучена с matryoshka-лоссом, поэтому первые 192 координаты 768-мерного эмбеддинга сохраняют почти всё качество. Наружу отдаются именно они, предварительно отнормированные. Вектор занимает в 4 раза меньше места в хранилище, и поиск кандидатов работает быстрее.

Настройка сервиса

Встроенный кэш включается в config.pbtxt модели директивой response_cache, а общий размер кэша задаётся при старте сервера параметром --cache-config. Кэш может быть локальным или на Redis. При профиле нагрузки эмбеддера запросов часто повторяются одни и те же входные данные, поэтому кэш снимает большую часть работы модели. Чтобы после рестарта поток трафика не обрушивался на непрогретый сервис, реализован прогрев: при старте пода прогоняются топ-N запросов, покрывающих большую часть трафика.

Батчинг настраивается через max_batch_size и dynamic_batching с preferred_batch_size и max_queue_delay_microseconds. Количество экземпляров модели задаётся в блоке instance_group, там же выбирается CPU или GPU. Для onnxruntime backend отдельно настраиваются intra_op_thread_count и inter_op_thread_count: первый определяет, сколько ядер получит один экземпляр модели, второй — потоки для независимых операторов графа.

Чернышев предупреждает: если суммарно потоков окажется больше, чем выделено или доступно, сервис начнёт тротлить — latency вырастет, а следом появятся ошибки по таймаутам на клиентах.