Mayorov BLOG

Спойлер: переписываю их на Rust. Нет, правда, дальше можно не читать — это всё, что я делаю.

Строго говоря, это не ускоряет сам инференс, но главное — ускоряет ответ сервиса пользователю. Причём значительно: где-то в 15-20 раз. То же самое можно провернуть и на Go, но Rust кажется более ML-friendly, плюс у нашей команды в нём большая экспертиза — мы писали на нём как различные движки симуляций, так и веб-сервисы (правда, без машинного обучения).

Наша архитектура

У нас сейчас выстроена комбинированная система для инференса, которая состоит из таких компонентов:

Triton Inference Server — индустриальный стандарт для инференса. Используем его в основном в связке с Airflow для моделей, которые требуют автоматического дообучения по расписанию или при необходимости (дрейф данных, деградация качества и др).

FastAPI Middleware для препроцессинга и сложной бизнес-логики при обработке и подготовке данных для инференса или преобразовании его результата. Если вы когда-нибудь писали кастомный бэкенд в Triton, то понимаете, почему проще написать свой middleware на FastAPI для взаимодействия.

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

Звучит просто? На деле это означает, что для банальной аугментации или нормализации данных вам нужно:

1. Создать Python-бэкенд для препроцессинга, который наследуется от triton_python_model.TritonPythonModel.
2. Написать конфиг для этой «модели» препроцессинга с описанием входов/выходов.
3. Настроить ансамбль, где явно прописать граф выполнения: препроцессинг → основная модель → постпроцессинг.
4. Далее — дебажить сериализацию между моделями: Triton передаёт данные через свой внутренний формат, и любая ошибка в типах или размерностях приводит к невразумительным ошибкам.
5. Мучиться с деплоем — каждая модель в ансамбле должна быть отдельно версионирована и загружена.

А теперь представьте, что вам нужно добавить условную логику: «если тип запроса А — делать препроцессинг X, если Б — препроцессинг Y». В обычном веб‑сервисе это if‑statement, а в Triton это отдельные ансамбли или костыльная логика внутри Python‑бэкенда, которая всё равно должна возвращать фиксированные типы.

Просто контринтуитивная штука, которая геометрически усложняется с увеличением требуемого числа шагов.

Отдельные FastAPI ML-сервисы — они нужны, когда мы хотим задеплоить модель с большим количеством кастомной обработки (например, фоновое дообучение) и нагрузка на коммуникацию превышает нагрузку на инференс. В нашем шаблоне уже есть поддержка gRPC, OpenAPI, кэширование, телеметрия и версионирование моделей.

Чистый Triton с Airflow мы используем в 10-20% сервисов (в основном LLM, потому что для их деплоя NVIDIA уже всё подготовила), такой же процент — для отдельных сервисов, и все остальные сервисы — это связка FastAPI middleware и Triton.

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

Посмотрели метрики, расстроились из-за скорости ответа и определили набор проблем, которые требуют решения:

1. Были неоптимальные предобработки, которые почему-то пропустили на ревью
2. Некоторые сервисы до сих пор сидели на старом шаблоне с pandas (мы уже давно перешли на polars)
3. Некоторые сервисы до сих пор использовали старый шаблон с корявым OpenAPI (слишком общая типизация)

Я решил поработать с cProfile и выделить узкие места в каждом Python-сервисе, похрустел пальцами, сконцентрировался... ну и устал уже через час. Еще через 5 минут я просто написал в чат: «Блин, давайте я просто перепишу это херню на Rust?»

В этот момент я понял, что воплощаю программерские мемы в жизнь

«Если всё значительно ускорится, я напишу новый Axum-шаблон, и мы впредь будем использовать его в критичных к скорости сервисах».

Чтобы питонисты не взвыли, уточнил, что от FastAPI мы никуда не уходим 😅

Опыт переписывания

В общем, переписать оказалось значительно быстрее, чем рефакторить этот python-слоп. Основной момент, который хотел бы выделить: дообучение в Rust лучше не делать — лучше выделить его в отдельный сервис на FastAPI, а полученный артефакт загружать в S3 и подтягивать в Rust для инференса.

Дообучение моделей прямо в Rust доставило очень много головной боли. ONNX Runtime вёл себя очень странно — побороть необъяснимые сбои дообучения мне не удалось, хотя буквально та же самая логика в FastAPI работала прекрасно.

Результаты

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

К сожалению, у меня нет скриншотов метрик, но у меня есть кое-что получше — честное слово 😉 С переходом на Rust работа всех бэкендов драматически ускорилась, особенно впечатлил сервис, запрос на predict в котором стал занимать вместо 10 секунд - 0,1 секунду (x100 speed boost получается).