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

Ссылка на проект на Github для особо нетерпеливых

Введение

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

По итогу мы имеем the next big thing, которая, в отличие от углеродных нанотрубок и Web3.0, не собирается потухнуть за год-два (на очереди еще квантовые компьютеры, но черед их хайпа еще не пришел). Ситуация на глобальном рынке похожа на ту, что мы имели во время открытия Нового Света: резко появляется новая, ранее неведомая область, с помощью которой можно делать вещи, о которых можно только было мечтать. С одной стороны, все компании подряд суетятся, пытаясь встроить модели в свои процессы, заняв место под Солнцем, и пользователи, с другой стороны, закономерно быстро привыкают к хорошему и уже ожидают наличие подобного функционала везде и всюду. Выходит замкнутый круг, в котором приходится в разгорающийся костер моделей машинного обучения подбрасывать всё новые и новые модели и процессы, связанные с ними, чтобы тот не затух.

Можно еще долго анализировать значение и роль ML в наше время, но смысл данной статьи кроется в очевидных вещах:

Ситуация также похожа на рассвет ReactJS, когда он также относительно резко стал популярным, но многие разработчики явно были не готовы писать в новом стиле компонентов и реактивности (которые существовали и до реакта, но не имели столь широкую популярность), особенно те, кто перешел во фронтенд недавно или не имел опыта создания production ready приложений с нуля. Всё же hello world примеры имеют мало общего с большими проектами. В итоге через пару лет мы имели самые разнообразные структуры React проектов, не было единых стандартов и прочего, из-за чего каждый городил как видит или как получается. Приправив это все тем, что многие разработчики приходят в сферу не для того, чтобы неспешно смакуя процесс разработки создавать налучшую и наиболее гибкую архитектуру, а для того, чтобы как можно проще и быстрее сделать то, что хочет заказчик, мы получаем огромное множество вариаций кодовых баз, большинство из которых едва ли можно назвать примером для подражания. Похожие вещи закономерно случаются каждый раз, когда что-то хайповое входит в нашу жизнь резко (будь то ООП или PHP): к этому никто попросту не готов. И лишь по прошествии многих лет приходит понимание того, что произошло, как делать не надо, а дальнейшая рефлексия и анализ преподносят нам заветные архитектурные паттерны и how to гайды. И, если пользователи не особенно страдают от этого и редко замечают взросление технологии, то мы сами, как разработчики, непосредственно варимся в этом котле попыток сделать так, чтобы потом не надо было постоянно хвататься за голову и избежать частых рефакторингов в дальнейшем.

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

Анализ существующих сервисов

Как пример такого приложения можно рассмотреть нашумевший пол года назад сервис BaiRBIE.me, предлагавший генерацию изображения того, как бы вы выглядели, если бы были Барби или Кеном. Хайп сервиса можно объяснить не только тем, что он вышел в то же время, что и ожидаемый фильм года - Barbie, но и тем, что подобного функционала раньше просто не было. Сам сервис представлял из себя простой веб-интерфейс с полем для загрузки фото и формочкой для указания предпочтительных параметров для генерации.

А теперь представим, что мы - один или несколько разработчиков. И в какой-то момент мы решили, что вскоре выходит ожидаемый фильм и его ожидаемость может нам сыграть на руку, так как привлечет пользователей. Итак, мы имеем относительно небольшие ресурсы в плане:

• разработчиков

• времени

• серверных мощностей

Однако в надежде на успех такого сервиса мы должны понимать, что чем больше кол-во пользователей, тем больше нагрузка на наш сервис, а масштабирование не выглядит весьма простым и очевидным, особенно когда речь заходит про использование ML в продакшене. Тем более, если мы говорим об автоматизации масштабирования и балансировки нагрузки. И данная ситуация, когда мы не можем себе позволить кучу специалистов ML и DevOps и надеемся, что в наш сервис могут за несколько дней прийти тысячи и миллионы пользователей, ставит нас в положение между молотом и наковальней.

Highload за маленькие деньги - мечта любого заказчика, что тут сказать :-)

Наш пример

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

В качестве интерфейса взаимодействия у меня будет обычная веб-страничка, как в случае BaiRBIE.me, т.к. нам понадобиться лишь демка, чтобы потыкать. Для реального приложения понадобится Rest API с webhook или очереди, в зависимости от требований.

В качестве txt2img модели я выбрал опенсорсную Stable Diffusion. Т.к. нам требуется фотореалистичный результат, мною был выбран чекпоинт epiCPhotoGasm-last-unicorn (не уверен, можно ли здесь публиковать ссылку на модель с civit.ai, но если надо - можно найти там и посмотреть примеры результатов). В качестве либы для использования модели (т.к. в нашем случае чекпоинт - это просто веса модели) я взял diffusers. Метод сэмплирования/scheduler - DPM++ 2M Karras (ссылки не нашлось, но уверен рядовые пользователи AUTOMATIC1111 webui его знают, собственно вот обсуждение того, как применить данный сэмплер с помощью diffusers.schedulers). Также воспользовался negative embedding - EasyNegative (также можно найти на civit.ai) для улучшения качества генерации. Остальные настройки более мелкие, они разбросаны по сущностям схемы таски на генерацию и самим классом генератора.

Способы реализации основной архитектуры

Монолит

Первое, что может прийти в голову

Ну че там сложного: принимаешь запрос на генерацию, генерируешь и выдаешь результат. Чего мудрить то?!

Такая реализация едва ли справится с большим наплывом пользователей по следующим причинам:

• т.к. процесс генерации является синхронным, в идеале занимая 100% ресурса GPU, то распараллелить данный процесс в рамках одного GPU будет невозможно. Из-за чего придется сохранять HTTP соединение с клиентом в течение всего процесса ожидания своей очереди на генерацию и самой генерации. Чем больше висящих соединений, тем больше нагрузка и на сеть, и на сервер

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

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

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

Модули, библиотеки, FFI и запуск подпроцесса с пайпами

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

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

Вариант с запуском подпроцессов в данном случае выглядит как изобретение велосипеда, если за велосипед брать способы, описанные далее: в лучшем случае вы создадите весьма упрощенный аналог Docker/systemd для всех компонентов вашего сервиса, но, вероятнее всего, загнетесь на каком-то этапе. Стоит ли игра свеч?

Разделение на микросервисы

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

Идиоматически достаточно разделить всего на 2 микросервиса - сервер и воркер:

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

Способы сетевого межпроцессного взаимодействия

HTTP

Опять же, это первое, что приходит в голову, когда речь заходит о взаимодействии между разными процесами, особенно, когда эти процессы находятся на разных машинах. Однако для нашего случая вариант не подходящий по причине необходимости удерживания соединения и потока/таски для каждого ожидающего пользователя. Больно представить, что случится, если в такую систему придет хотя бы 100 пользователей. И ведь им даже не обязательно отправлять запрос одновременно: достаточно делать каждый новый запрос с интервалом меньше времени генерации одного результата, чтобы таски начали накапливаться. А там и до также вышеупомянутого таймаута недалеко, вот пользаки будут рады)

RPC, gRPC

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

Message queues

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

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

В качестве брокера очередей я буду использовать RabbitMQ.

Для взаимодействия между сервером и воркером мы будем использовать 2 очереди:

1. для отправки тасок от сервера к воркеру на выполнение

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

Но как нам передавать результат генерации? Ведь, в зависимости от модели и прочих параметров, это может быть текст на 1КБ, а могут быть и широкоформатные изображения или даже видео на многие мегабайты - заваливать очередь столь тяжелыми сообщениями не стоит. Кто знает, какой контент мы сможем генерировать в будущем? Давайте рассмотрим возможность передачи именно для ситуации большого контента, т.к. с маленьким контеном и так всё понятно. Тем более, что даже в нашем случае результат - это изображение.

Передача результата от воркера к серверу и клиенту

Файловая система

Опять первый вариант, пришедший на ум? И опять не подходящий, да?

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

Cloud object storage

Специальный сервис для хранения файловых объектов с возможностью доступа к данным из разных машин и распределения хранения? Выглядит неплохо!

В нашем случае я выбрал S3-совместимый сервис - MinIO. Сохраняем результат воркером в MinIO и далее можем получать доступ к нему с сервера.

Также можно включить процесс автоматической отчистки устаревших объектов, чтобы не нужно было делать это самому (в случае fs пришлось бы).

В итоге выходит следующая картина:

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

Ожидание результата от сервера

Polling/Long polling

Очевидный вариант, который отстой, да?

Да, но в случае нашего приложения я применил именно простой polling

А всё потому, что в нашем случае совершенно необязательно бомбить сервер запросами вроде "уже готово?", достаточно опрашивать его раз в некоторый интервал. В моем случае процесс генерации изображения 512/512 занимал примерно 2.5 секунды, так что с точки зрения пользовательского опыта отправлять запросы чаще данного интервала необязательно - если пользователь спокойно готов подождать этот период времени и уж тем более нескольких человек перед ним, то можно опрашивать сервер раз в 3-5 секунд.

Плюс данного метода состоит в том, что таким образом нагрузка на сервер снижается значительным образом!

Забегая вперед, отмечу: на веб-странице для создания динамики я воспользовался хайповым HTMX и Materialize CSS, т.к. тащить реакт или что-то похожее в простейший веб-интерфейс не хочется, особенно учитывая факт того, что потом надо будет нагружать сеть еще и отдачей статики.

WebSockets/SSE

Данные протоколы, в отличие от stateless HTTP, являются stateful, то есть хранят состояние.

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

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

Оптимизация polling

Рассмотрим процесс запроса, которым будут спамить веб-страницы клиентов при polling:

1. Запрос приходит на сервер

2. Сервер идет в MinIO, чтобы проверить, не появился ли там результат генерации

3. Сервер отдает ответ клиенту

Как можно догадаться, запрос в MinIO в данном случае можно оптимизировать. Как? Просто заменив его.

Заменить данную проверку можно воспользовавшись любой in-memory базой данных, т.к. та хранит данные в ОЗУ, тем самым ускорив процесс проверки на наличие результата генерации. В общем, обработку запроса можно представить следующим упрощенным псевдокодом:

function get_result(id):
  is_ready = in_memory_db.exists(id)
  if !is_ready:
      return 404

  return minio.get(id)
end

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

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

Масштабирование

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

Сервер

В моем случае сервер был реализован на языке Go. И с точки зрения идеологии языка, и с точки зрения простоты нашего API, я не стал использовать сторонние бекенд фреймворки, а воспользовался встроенным модулем net/http. Для отдачи статики и генерации кусочков HTML, т.к. мы используем HTMX - также встроенный модуль text/template.

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

Способ не является чем-то новым и также имеет ограничения: процессорное время и память не бесконечные. Хотя сами горутины являются весьма эффективными и вероятно само приложение смогло бы обрабатывать запросы от тысяч одновременных пользователей, масштабирование приложения можно проводить путем поднятия дополнительных инстансов на других машинах и связывании их одним прокси-сервером и балансировщиком нагрузки. Из готовых решений вспоминается достаточно мощный веб-сервер - Nginx. Traefik также может быть использован в силу своей нацеленности на маршрутизацию между контейнерами, в которых как раз развернуто наше приложение, и из-за его интеграции с Let's Encrypt, во избежание необходимости настройки связки Nginx и certbot для менеджмента TLS-сертификатами для запуска приложения под HTTPS.

Воркер, ускорение генерации

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

Ну тут же всё очевидно - сервак на гошке написан, а воркер - на питончиге! Вот вам и вся разница!

Дело едва ли в этом. Особенно, учитывая факт того, что практически всё время занимает именно процесс генерации библиотечкой diffusers, использующей под капотом PyTorch, использующий C/C++ и CUDA для тензорных вычислений.

Выходит, что для ускорения обработки воркером выполнения тасок необходимо как-то ускорить процесс самой генерации (спасибо, дядя Кэп).

Во-первых, стоит использовать GPU на 100% при выполнении генерации одного изображения. Т.к. даже при параллелизации всё равно будет выходить одно и то же общее время, а при выполнении генерации одного изображения - чем больше CUDA-ядер будете использовать, тем быстрее будет выполняться процесс генерации. Совет, скорее, для более низкого уровня, т.к. в моем случае diffusers используем GPU на максимум по умолчанию.

Сам воркер использует синхронный Gunicorn, ранящий ассинхронные Uvicorn воркеры. Вспоминая всеми любимый GIL, мы приходим к тому, что с помощью гуникорна, использующего многопроцессорность по умолчанию, можно кратно кол-ву воркеров распараллелить задачу. Однако тут не всё так просто! И следующий тест продемонстрирует это:

Я запускал приложение с 1-4 воркерами воркера (простите за тафтологию), при этом кратно росло лишь кол-во потребляемой RAM и VRAM, и, конечно же, процессорной мощности. Результат представлен ниже:

Дело в том, что таким образом выйдет лишь увеличить прием сообщений из очереди и отдачу сообщений в очередь обратно. Сам процесс будет занимать столько же времени, т.к. все инстансы будут делить ресурс GPU между собой.

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

Дальнейшее развитие проекта

• Мониторинг: можно развернуть систему вроде Prometheus с Grafana и Portainer. Healthcheckи в моем примере уже есть и их можно успешно заюзать в них.

• Оркестрация: проект уже использует Docker Compose, но более продвинутые фанаты, конечно, могут заюзать Kubernetes

• Система краш-логов: Sentry

...мб еще что-то, напишите в комментариях, чего явно не хватает

Заключение

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

Еще раз продублирую ссылку на Github проекта.

Буду рад замечаниям по статье и обсуждениям!