Это третья статья по теме реализации масштабируемой системы для выполнения задач распределенного машинного обучения на GPU с использованием Java, Kotlin, Spring и Spark. Список всех статей: Варианты использования Java ML библиотек совместно со Spring, Docker, Spark, Rapids, CUDA Масштабируемая Big Data система в Kubernetes с использованием Spark и Cassandra Использование Kotlin и WebFlux для выполнения задач ML в Apache Spark на GPU В предыдущей статье для создания Spark Driver приложения использовался сервлетный стек Spring (Boot 2.7.11) и JDK 8.  На дворе вторая половина 2023 года, у многих в проде уже используется Boot 3+ (а то и 3.1+), совсем скоро должна выйти новая LTS версия Java, и, мягко говоря, Boot 2+ и JDK8 устарели. Использовались они намеренно, так как для задач тренировки моделей машинного обучения на GPU в среде Spark частью системы является ускоритель вычислений на GPU NVidia Rapids. Поддержка JDK 17 появилась только в релизе v23.06.0 от 27.06.23, с ее выходом появилась возможность перейти на актуальную LTS версию Java, а с ней - на Spring Boot 3+. В данной статье описывается миграция с Boot 2 и JDK 8 До Boot 3 и JDK 17, со Spring Web на Spring WebFlux, в конце сравниваются Web и WebFlux версии по потреблению аппаратных ресурсов и скорости выполнения. Для миграции достаточно поднять версии Rapids до 23.06.0, JDK до 17, Spring Boot до 3.1.1. Нюансов не так уж и много: Конфликт логеров Slf4j и Log4j при использовании Spark: из зависимости spring boot starter web исключаем spring boot starter logging:  Запускать Spark Driver на JDK 17 необходимо со следующими параметрами (приведено для Dockerfile):  В связи с переходом на Hibernate 6, при использовании JSOB и BYTEA полей в сущностях Postgres придется немного отрефакторить Entity: При этом, использовавшийся ранее CustomPostgresDialect оказывается не нужным и его можно удалить, заменив на org.hibernate.dialect.PostgreSQLDialect:   Не считая докерфайлов и действий по добавлению новой версии jar'ника Rapids в директорию с jar-файлами для отправки в Spark executors и в образ executor’а, это все, что необходимо выполнить. Актуальную версию можно взять в соответствующей ветке репозитория. На этом можно было бы и закончить, но любопытство ведь берет свое, и появился вопрос - а заработает ли на реактивном стеке, и будет ли эффект? Изменений при таком переходе изначально должно быть больше, но так же есть нюансы в виде управления зависимостями. Так, Netty, необходимый для Project Reactor (WebFlux) используется самим Spark и драйвером Cassandra, поэтому изначально конфликтовали. Решается путем задания трех зависимостей в самом начале списка зависимостей:  Spring Data тоже заменяется на реактивную версию:  И добавляются несколько библиотек для работы Kotlin в среде WebFlux:  Кстати, сам Kotlin тоже поднял с версии 1.8.21 до 1.9.0. Для логирования HTTP запросов-ответов добавляем Zalando Logbook:   Модифицируем Main класс приложения, необходимо добавить аннотации @EnableWebFlux и @EnableR2dbcRepositories, указать тип приложения REACTIVE  Так как в сущности БД используется JSONB поле (с его отображением в приложении в виде JsonNode), необходима конфигурация R2DBC с кастомными конвертерами:  Далее следует удалить из упомянутой выше сущности ModelEntity лишние аннотации, в итоге должно получиться:  Сам репозиторий сущностей теперь наследуется от R2dbcRepository: Методы сохранения и загрузки модели трансформируются для работы в WebFlux: метод загрузки модели из БД модифицируется до: а метод сохранения: модифицируется до:  Репозитории Cassandra строились на основе взаимодействия со спарковой сессией. Переработать методы довольно просто. Так, метод получения датасета в базовом абстрактном репозитории:  модифицируется до:  метод сохранения датасета:  модифицируется до:  метод получения базового датасета с определенными оффсетами:  модифицируется до:  Остальные репозитории конкретных таблиц переписываются по такому же принципу. В сервисе работы с данными следует упомянуть метод объединения датасетов (теперь же репозитории возвращают реактивные Mono<Dataset<Row>>):  модифицируется до:  здесь получаются 8 датасетов в Mono-обертках, обертки объединяются в один Mono посредством .zip() и передаются на исполнение в метод комбинации датасетов, который не менялся. Метод выполнения предикта с помощью сохраненной модели:  модифицируется до:  Метод обучения модели:  модифицируется до:  здесь tdf и edf обернуты в Mono, поэтому объединяются в кортеж из двух элементов Mono<Tuple2>, далее в оборачиваем в Callable функцию regressor.fit(tdf), которая будет выполнена асинхронно и вернет результат в виде model: XGBoostRegressionModel. В функции flatMap она используется с эвалюирующим датасетом для получения предиктов, затем сохраняется в БД с помощью описанного выше метода saveModel. Остальная логика очевидна. Наибольшую сложность вызывает метод инкрементального обучения (да, на данной модели инкремент не работает и требуется замена XGBoost на другую модель, но цель была трансформировать логику под реактивную среду и получить работающий пример, который далее можно использовать для инкрементального обучения модели).  Исходный метод:  модифицируется до:  В отличие от Java, лямбда-выражения в Kotlin не требуют от переменных, чтобы они были effectively final, поэтому переменные currentBatchOffset, i, model и tdf могут изменяться в ходе выполнения основного стрима. Здесь функция получения датасета обертывается в Mono.defer(). Особенность данного подхода в том, что выполнение функции откладывается до момента подписки на данный Mono. А подписка будет повторяться методом .repeat() до тех пор, пока не выполнится условие tdf?.isEmpty == false. Когда очередной tdf будет пустым, выполнится логика в then: из кассандры будет получен датасет edf, который сохранится в таблице скомбинированных данных, так же будут получены предикты модели и сохранена сама модель. Затем из предиктов подготовится результат метода. В случае ошибки вернется пустой результат метода. Не сказать, что это идеальное исполнение метода, но как пример сойдет. Подробно реализацию можно посмотреть в отдельной ветке репозитория. Как известно, реактивный стек отличается от сервлетного тем, что для выполнения одной и той же логики зачастую требуется меньше ресурсов. В некоторых случаях может возрасти скорость выполнения алгоритма. Тестирование происходило по следующей методике:  Сервис поднимается в Docker-контейнере с 4 CPU и 4 Gb памяти, использует Spark Executor (v. 3.3.2, JDK 17), так же в Docker контейнере, который подключается к Standalone-мастеру Spark в виртуальной машине. Все работает на одной машине под управлением Windows 10 Pro, задачи тренировки моделей выполняются на GPU NVidia 4090. В течении 10 минут производятся запросы методов: обучения новой модели (POST /analytics - для сокращения “1 запрос”), получения предиктов с помощью сохраненной модели (GET /analytics - для сокращения “2 запрос”) и инкрементального обучения (POST /analytics/increment - для сокращения “3 запрос”) с batch_size = 50 записей, во время которого делается 12 итераций над 6 сотнями записей в таблицах Cassandra. Первый цикл на “не прогретом” драйвере (первые запросы всегда выполняются дольше), далее два одинаковых цикла по одному запросу каждого метода на “прогретом драйвере” и в четвертом цикле запускаются одновременно 1, 2, 3 методы. Driver работает в режиме Spark Cluster, используется одна Spark Session на все время работы приложения; Изначальные параметры запуска JVM одинаковые: первоначальный размер кучи 512 Мб, максимальный размер не указан, GC по умолчанию (G1). Результаты потребления ресурсов:  Максимальное потребление CPU Среднее потребление CPU Максимальное потребление памяти, Gb Среднее потребление памяти, Gb Количество Stop the world за 10 минут Spring Web 3,4 1,5 4 2 4 Spring Webflux 3,4 1,1 1 0,5 0 С указанными выше параметрами для сервлетного стека наблюдалось 4 stop the world от G1 GC, при этом один раз результатом выполнения предиктов из сохраненной модели стала ошибка сервера. На графике видно, что потребление памяти растет линейно до момента, когда свободного места для кучи уже нет и ее необходимо чистить. У реактивного стека другая картина: после первых запросов стабильные ~0.5 Гб памяти. По потреблению CPU разница не настолько большая. Скорость выполнения запросов: Топ-5 классов по потреблению памяти:   Учитывая, что весь результирующий датасет занимает около 55 Мб, такой объем аллоцированной памяти вызывает вопросы. Анализ стектрейсов показал, что в большинстве случаев источником и причиной является Spark и Rapids, которые строят план запросов, обмениваются данными между БД, экзекуторами и драйвером, подготавливают массивы данных для загрузки в GPU и вычитывают результат из него. Потратив некоторое время на изучение вопроса оптимизации использования памяти, могу сделать вывод, что это штатное поведение системы в такой конфигурации, и надо научиться с этим жить при использовании сервлетного стека. Первые попытки жить с этим в привели к изменению параметров запуска JVM для сервлетного стека на следующие: -Xms512m -Xmx3g -XX:GCTimeRatio=19 (жесткое указание того, что система может потратить до 5% времени на сборку мусора - (1 / (1+19))) -XX:+UseZGC. Учитывая, что реактивному стеку достаточно в среднем 512 Мб памяти, и что Z GC потребляет несколько больше памяти, чем G1 GC, планка максимального размера кучи снизилась до 3 Гб.  Потребление CPU незначительно снизилось, наблюдается схожее потребление памяти, но stop the world уже не зафиксировано. Судя по графику, после завершения работы методов POST /analytics и GET /analytics куча очищается, но при работе POST /analytics/increment куча очищается только к моменту приближения к своему максимальному размеру. Логика, которая могла бы вести к утечке памяти, отсутствует, причина такого высокого потребления памяти остается не выясненной.  Результаты переключения GC в таблице потребления ресурсов:  Максимальное потребление CPU Среднее потребление CPU Максимальное потребление памяти, Gb Среднее потребление памяти, Gb Количество Stop the world Spring Web 3,4 1,5 3 1,5 0 Spring Webflux 3,4 1,1 1 0,5 0 и скорости выполнения запросов: Стало интересно, что будет, если для G1 GC установить максимальный размер кучи как для Z GC и установить жесткий предел времени выполнения на сборку мусора. В этом случае оказалось, что память заполняется как и раньше, но stop the world стало больше, так как доступной памяти меньше, и, соответственно, заполняется она быстрее. Потребление ресурсов осталось примерно на том же уровне:  Максимальное потребление CPU Среднее потребление CPU Максимальное потребление памяти, Gb Среднее потребление памяти, Gb Количество Stop the world Spring Web 3,4 1,5 3 1,5 6 Spring Webflux 3,4 1,1 1 0,5 0 Скорость выполнения запросов возросла, но не существенно. Так же попробовал в сервлетном стеке использовать Parallel GC с параметрами -Xms512m -Xmx4g -XX:GCTimeRatio=19 -XX:+UseParallelGC. Результаты самые худшие, за 10 минут удалось прогнать только 2 цикла. Если первые два метода выполнялись примерно за то же время без отклонений, то метод инкрементального обучения выполнялся в первый раз 3мин 32с, что хуже примерно на 1,5 минуты среднего результата сервлетного стека, а второй запрос подвис и выполнялся 8мин 10с. Результаты в таблицах не фиксировал. Напоследок применил настройки JVM -Xms512m -Xmx3g -XX:GCTimeRatio=19 -XX:+UseZGC к версии WebFlux, которая оказалась на данный момент самой оптимальной по потреблению ресурсов и скорости обработки запросов. Сравнительные таблицы версии с дефолтными параметрами и G1 GC и версии с кастомными параметрами JVM с Z GC ниже. График потребления ресурсов: Таблица потребления ресурсов:  Максимальное потребление CPU Среднее потребление CPU Максимальное потребление памяти, Gb Среднее потребление памяти, Gb Количество Stop the world Spring Webflux G1 GC 3,4 1,1 1 0,5 0 Spring Webflux Z GC 3,4 1,4 2,84 1,5 0 Таблица скорости выполнения запросов:   Итоговые таблицы со всеми версиями и различными значениями параметров конфигурации JVM представлены ниже. За основу взяты результаты для WebFlux на G1 GC с одним параметром JVM минимального размера хипа 512m. Подводя черту после написания третьей статьи на тему построения системы распределенного машинного обучения на Java и Kotlin, самый большой вывод, который напрашивается - построить подобную систему сложно, много неизвестных, необходимо выполнить много исследований, но добиться работающего решения вполне реально, было бы желание. Если так случилось, что нужно выполнять задачи ML на JVM стеке технологий, учите Python и не занимайтесь фигней, а руководству продайте альтернативную систему отличным выбором в качестве основы будет Kotlin и Spring Webflux (как альтернатива - Web с Z GC), и, естественно, Apache Spark. По окончанию работ над любым приложением стоит провести проверку профилировщиком, так как с очень высокой вероятностью при дефолтных параметрах JVM работа приложения не будет оптимальной. Другой вопрос - является ли данная система эффективной с точки зрения производительности и потребления ресурсов? Без тестов на альтернативной системе (например, Python + Dask) объективно ответить на данный вопрос я затрудняюсь. Возможно, в будущем попробую поднять такую систему и написать альтернативную логику на питоне, тогда будет с чем сравнить и о чем написать очередную статью.