14.01.2025 08:05
ledevik 0 594 Источник

Системный разработчик ИТ-компании «Криптонит» написал статью про новый инструмент на Rust, который облегчает запуск моделей машинного обучения и их внедрение в приложения. Дальше публикуем текст от первого лица.

Привет, Хабр! Меня зовут Михаил Михайлов. Я пишу на Rust и работаю системным разработчиком в компании «Криптонит». В этой статье хочу рассказать о новом инструменте на Rust, который облегчает запуск моделей машинного обучения и их внедрение в приложения.

Это эффективная и универсальная библиотека (в языке Rust называемая crate или крейт) с открытым исходным кодом, которую мы изначально написали для своих нужд. С её помощью вы сможете запустить практически любую готовую модель ML!

Поддерживаются не только профессиональные ускорители Nvidia, но и  широкий спектр потребительских видеокарт разных поколений (Maxwell, Pascal, Volta, Turing, Ampere, Ada Lovlace).

Просто добавь ИИ!

Машинное обучение повсюду: большие языковые модели, генеративные диффузионки, рекомендательные системы... Сейчас сложно найти предметную область, в которую кто-то не засунул «машинку».

Однако, несмотря на широкие возможности ML, работа с ним часто кажется слишком сложной. Мы создали инструмент, который облегчит работу с ML для разработчиков на Rust. С ним вы увидите, что этот процесс не только легко доступен, но и наделяет ваши приложения невероятной мощью ИИ!

Представьте разработчика, который не разбирается в тонкостях ML, но может скачать готовую модель. Допустим, ему надо не просто встроить чужую модель в свой легковесный проект, а интегрировать её в амбициозно высоконагруженный конвейер (pipeline). При этом разработчик не готов страдать — он же не DevOps! Ему хочется получить простой и безболезненный деплоймент. Исходя из этих условий, мы выделили следующие требования для идеального инструмента:

  • поддержка разных форматов моделей;

  • совместимость с широким спектром видеокарт;

  • универсальность и эффективность;

  • лаконичность и гибкость.

Популярные решения для интеграции моделей ML

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

Например, сначала мы запускали модели при помощи обёрток (bindings) к движкам TensorRT и CUDA. Однако данный способ требовал компилировать каждую модель под конкретную тройку версий TensorRT, CUDA и модели GPU.

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

Другой популярный вариант — Deepstream SDK. Это продукт Nvidia с закрытым исходным кодом. Он позволяет крайне эффективно исполнять конвейеры машинного зрения (CV pipelines). Платой за производительность является его универсальность:

  • почти невозможно вносить изменения в исходный код;

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

  • работает только с видео. А мы хотим что-то всеядное, что готово переварить и картинки, и аудиозаписи… почти любые модели.

Довольно долго мы использовали инфраструктуру с отдельно поднятым сервером, на котором исполнялись модели. В таких условиях мы опробовали Triton Inference Server и TensorFlow Serving. Второй нам показался похуже в плане универсальности и эффективности, поэтому мы почти сразу его отбросили.

Некоторое время нас в целом устраивал Triton Inference Server: нравилась его гибкость, универсальность и надёжность. Однако по мере роста нагрузки мы столкнулись с тем, что передача данных происходит неэффективно. Она идёт по сети, что порождает две проблемы:

Первая заключалась в том, что производительность конвейера упиралась в пропускную способность внутренней сети. Вторая была глубже: в некоторых сценариях приходилось копировать данные из GPU и пересылать их по сети, чтобы затем снова скопировать в GPU. Это съедало большой процент ресурсов и снижало общую производительность.

Когда нам надоело с этим мириться, мы перешли к созданию собственного крейта — Tritonserver-rs. Было решено взять всё самое лучшее от Triton Inference Server, отказавшись от отдельно вынесенного сервера.

Собственная библиотека на Rust

В итоге был написан крейт, основанный на обёртках к библиотеке Triton на Си (и в меньшей степени к сишным библиотекам CUDA).  Он обеспечивает асинхронное взаимодействие типа «запрос-ответ» с моделями внутри локально поднятого Triton Inference Server. По сути это библиотека с открытым исходным кодом, у которой под капотом Triton C-lib и пара фишек.

Наш инструмент позволяет встраивать модели в Rust-приложения, добиваться их высокой производительности и эффективно справляться с растущей нагрузкой. Это универсальное решение, которое поддерживает разные форматы моделей (TensorFlow, ONNX, PyTorch, OpenVINO, кастомный формат и пайплайны моделей).

Tritonserver-rs универсален и особенно удобен тем, кто занимается построением систем, использующих несколько моделей ML. Например, мы применяем пайплайны для систем распознавания и транскрибации речи, трекинга объектов на видео и их классификации, а также других задач, типичных для ИИ.

Из-за того, что под капотом нашего крейта лежит Triton Inference Server, мы наследуем все его плюсы:

  • удобный формат конфигурирования моделей;

  • единый запуск для всех архитектур, GPU и форматов моделей из коробки;

  • работа с несколькими моделями ML одновременно;

  • поддержка Prometheus.

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

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

  • пользовательское управление памятью, в том числе возможность отправлять запросы с GPU-данными;

  • высокоэффективная библиотека на Rust.

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

Конечно, не обошлось и без минусов. Как и 90% крутых штук в мире ML, наш инструмент заточен под аппаратное обеспечение Nvidia.

Развёртываем приложение с ИИ за несколько минут

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

Ниже я покажу как написать приложение на основе модели ML за 3 простых шага при помощи Tritonserver-rs. Давайте создадим программу, которая на вход берёт картинку, размечает на ней все объекты и классифицирует их. Под капотом будет стандартная для этой задачи модель YOLOv8.

Шаг 1

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

Стандартный конфиг модели выглядит так:

name: "yolov8"
platform: "onnxruntime_onnx"
max_batch_size: 8

input [{
   name: "IMAGE"
   data_type: TYPE_FP32
   dims: [3, 1280, 720]
}]

output [{
   name: "PREDICTIONS"
   data_type: TYPE_FP32
   dims: [256]
}]

В нём указываются имя модели, её формат, желаемые оптимизации, а также все входы и выходы модели. Последнее — это вектор структур, содержащих следующие параметры: имя входного тензора, тип данных, а также размерность этого тензора.

На этом сложности заканчиваются.

Шаг 2: пишем код

Добавляем crate в Cargo.toml:

[dependencies]
tritonserver-rs = "0.1"

Поднимаем локальный сервер:

let mut opts = Options::new("/models/")?;
opts.backend_directory("/path/to/backends/")?
.gpu_metrics(true)?;

let server = Server::new(opts).await?;

Далее создаём входные данные (в нашем случае: открываем картинку)

let image = image::ImageReader::open("/data/traffic.jpg")?.decode()?;
let input = Buffer::from(image.as_flat_samples_u8().samples);

Создаём запрос и прикрепляем к нему вход.

let mut request = server.create_request("yolov8", 2)?;
request
.add_input("IMAGE", input)?
.add_default_allocator();

Асинхронно ожидаем ответа от модели.

let response = request.infer_async()?.await?;

let result = response.get_output("PREDICTIONS");

Шаг 3: развёртывание

Как и обещал, деплоймент с Tritonserver-rs чрезвычайно прост. Вот как он выглядит на примере docker-compose.yml:

my_app:
   image: nvcr.io/nvidia/tritonserver:24.08-py3
   volumes:
     - ./application:/bin/application
     - ./models:/models
     - ./traffic.jpg:/data/traffic.jpg
   entrypoint: ["/bin/application"]

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

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

При этом для написания приложения нам понадобились всего 2 вещи: готовая модель и базовые знания о её входах и выходах.

Need for speed!

Нами был проведён тест производительности крейта на примере видеоконвейера, включающего обнаружение и трекинг объектов. Мы сравнивали скорость Tritonserver-rs, Deepstream и Triton Inference server. Последний запускался в двух режимам. В первом данные передавались по grpc из приложения на Rust. Во втором указатели CUDA передавались напрямую в модель при помощи библиотеки на Python.

Как видите, наше решение существенно превосходит внешний Triton Server. Более того, мы не сильно отстаём по скорости от Deepstream, который считается одним из наиболее эффективных инструментов видеоаналитики, изначально спроектированным под эту задачу.

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

Вот ссылка на крейт, примеры и документацию. Пользуйтесь наздоровье! Если возникнут вопросы — пожалуйста, задавайте их в комментариях.

https://crates.io/crates/tritonserver-rs
https://github.com/3xMike/tritonserver-rs
https://github.com/3xMike/tritonserver-rs/tree/main/examples
https://docs.rs/tritonserver-rs

Мы в активном поиске новых людей в нашу команду — посмотреть вакансии можно на карьерном сайте «Криптонита».

Статья написана по материалам выступления Михаила на RustCon 2024. Посмотреть видеозапись доклада можно в VK Видео.

Комментарии (0)