Это первая статья по теме реализации масштабируемой системы для выполнения задач распределенного машинного обучения на GPU с использованием Java, Kotlin, Spring и Spark. Список всех статей: Варианты использования Java ML библиотек совместно со Spring, Docker, Spark, Rapids, CUDA Масштабируемая Big Data система в Kubernetes с использованием Spark и Cassandra Использование Kotlin и WebFlux для выполнения задач ML в Apache Spark на GPU В данной статье рассматривается способ использования GPU nVidia с технологией CUDA в Docker-контейнерах для распределенной тренировки моделей машинного обучения на нескольких машинах. Цель статьи - показать вариант использования Big Data Tool Apache Spark в Docker-контейнерах, совместно с акселератором GPU вычислений Rapids на устройствах nVidia CUDA, с применением библиотек DJL, Spark ML, XGBoost, в приложении Spring Boot на Java 8 (требование Rapids), на нескольких машинах под управлением ОС Windows 10 Pro для решения задачи тренировки моделей машинного обучения в распределенной системе. Те же контейнеры в дальнейшем можно использовать в Kubernetes. Важное условие, из которого вытекают все болезненные решения - все действия выполняются в ОС Windows 10 Pro. Далее в статье будет рассмотрен рассмотрен вариант использования WSL2 (Linux Subsystem for Windows), но прежде важное требование - новейшие версии Rapids не работают с видеокартами на архитектуре Pascal, т.е. для запуска примеров в среде Windows требуется видеокарта nVidia на архитектуре Turing (серии 1600, 2000) и выше. Под OC Linux карты на архитектуре Pascal работать будут, рекомендуется запускать примеры ниже на Ubuntu 20.04 (наверняка, Debian 10 так же будет работать), но не выше - требование Rapids. Еще одним важным условием является реализация всех примеров именно на Java. В мире Spark (вместе со Spark ML) более распространен вариант использования Scala. Scala я не знаю, и особого желания изучать нет, а вот желание изучить Spark и ML для собственного развития и решения широкого круга задач имеется. Учитывая, что Scala и Java равнозначны в среде Spark, в отличие от того же Python, и что на Java существует ряд библиотек ML, которые можно использовать совместно со Spark, и принимая во внимание уже имеющийся опыт работы с ним, решение попробовать реализовать несколько примеров не заставило себя ждать. Код статьи проверен на Windows 10 Pro, GeForce RTX 2060 и 1080 Ti (с последней на Windows не заработало), часть скриншотов сделано во время настройки второй машины с картой GeForce 1650. Предваряя вопрос читателя, почему бы все не сделать на той же Ubuntu 20.04, отвечаю: а) так каждый сможет, вы на винде попробуйте; б) нет технической возможности (Cloud не вариант - дорогие машины с GPU). Целевая схема запуска представлена на рисунке ниже: Т.е. имеется от двух до n нод, на каждой из которых от 1 до m GPU устройств, docker runtime с контейнером Spark worker, из которого доступны GPU. Hardware и software слои описываются схемой: Докер позволяет запускать множество контейнеров с различными приложениями:   Это подходит для задач распределенной тренировки моделей ML в инфраструктуре Apache Spark: в настоящей статье рассматривается пример запуска Standalone кластера Apache Spark с одним Master узлом, двумя Worker узлами на разных машинах, и Spring Boot Java 8 приложением с использованием библиотек DJL, Spark ML и XGBoost в отдельном контейнере (спойлер - заработало не все, и не заработает без машины с Linux). Интересным является возможность использования Embedded устройств nVidia для IoT-устройств. Весь нижеописанный код доступен в репозитории GitLab. Все нижеследующие действия выполняются на Windows 10 Pro. Важно выполнять именно на конфигурации не ниже Pro, и сборке Windows 10 Build 19043.1263 (21H1). WSL, Docker и CUDA будут установлены в рамках данной статьи. Рекомендуемая версия WSL 5.10.16.3+; Docker 19.03+. Рекомендуется установить Windows Terminal для открытия множества вкладок терминала: PowerShell, cmd, Ubuntu. Убедитесь, что версия nVidia CUDA не ниже 11.7. Драйвер, который содержит данную версию, на момент написания статьи имеет версию 516.40.  Проверить версию драйвера и CUDA можно, открыв Powershell (лучше сразу открывать от имени администратора, но это требование для будущих действий) и выполнив команду  Для того, чтобы использовать GPU в Docker-контейнерах, необходимо установить ПО от nVidia (см ниже), которое требует установки WSL2. Если на ПК пользователя WSL не установлена, то можно установить командой ниже: Если WSL уже установлена, лучше обновиться до последней версии и проверить версию Ubuntu, должна быть 2. Требуется перезагрузка. После перезапуска установится Ubuntu для Windows в отдельном окне По окончанию установки можно проверить версию WSL в Powershell  В случае, если версия Ubuntu 1, следует ее обновить Следует установить Docker Desktop, если еще не установлен. Если установлен, рекомендуется обновить. На момент установки на систему без Docker Desktop, моему выборы были представлены следующие настройки: Я оставил оба чекбокса активированными. По окончанию установки требуется перезагрузка. В настройках необходимо убедиться, что чекбокс “Use WSL 2 based engine” активирован. Apply & Restart. Можно проверить, что в wsl появились новые записи в списке Проверить работу Docker можно командой Установится локальный docker registry, который будет полезен в последующей работе.   Ключевым моментом является поддержка работы CUDA в докер-контейнерах, для этого у nVidia есть решение.  Выполнить следующие шаги в PowerShell под именем администратора: Теперь нужно проверить в отдельном окне PowerShell работу тестового контейнера nVidia с флагом benchmark: Если вывод содержит нечто подобное, все шаги выполнены корректно и можно продолжать работу. Если имели место ошибки, рекомендуется обратиться к страницам документации nVidia здесь и здесь для их решения. Учитывая, что для работы Rapids необходимо использовать Java 8, следующие шаги по подготовке всех необходимых Docker образов, а в последующем и самих приложений, будут выполнены исходя из данного требования. Первоначально необходим самый базовый образ. Ниже листинг Dockerfile. Используется базовый образ Ubuntu 20.04 с CUDA 11.7.0 из репозитория образов nVidia. Доступный образ с Ubuntu версии 22.04 не подошел по причине совместимости всех компонентов системного ПО, необходимого для запуска прикладного ПО. Используемым в образе JDK является openjdk8, что соответствует требованиям Rapids и не испытывает проблем с лицензионным соглашением Oracle JDK. Также в образ устанавливается набор приложений для дебага. Инструкция CMD не обязательна, но удобна для отладки. Стоит заметить, что первоначально используется базовый образ nVidia с пометкой “devel” - тестирование происходило именно на нем, чтобы исключить возможные ошибки, связанные с недостаточностью компонентов. При этом, имеется образ: отличие от devel - отсутствие “nvcc”.  Собирается образ командой: Обращаю внимание, что на моей локальной машине имеется контейнер с репозиторием образов Docker, мне удобно при работе с локальным кластером Kubernetes указывать в манифесте свои образа из localhost:5000, и загружать их, не используя внешние репозитории.   Запускается контейнер командой: Примечание: важным флагом является “--gpus”, которому передается значение “all” - благодаря данному флагу контейнеру доступны все ресурсы gpu локальной машины. Проверить работоспособность базового образа можно путем выполнения в контейнере двух команд: При использовании базового образа “runtime” nvcc --version выведет ошибку, так как nvcc отсутствует в данном образе. Если нет похожего вывода какой-либо из команд, следует вернуться к предыдущим разделам и проверить корректность выполнения всех шагов. Следующий шаг - подготовка образа Spark Worker. Здесь следует отметить, что в данной статье рассматривается запуск кластера Spark как Standalone кластера, без менеджера ресурсов. Spark Master запускается на локальной виртуальной машине (у меня уже имелся настроенный мастер на виртуальной машине для тестирования работы Spark с Cassandra в рамках другой задачи, которая в настоящей статье не рассматривается), и к ней подключается Spark Worker в Docker контейнере. Полезность данного теста состоит в том, что: а) тестируется работоспособность GPU-нагрузки в контейнерах; б) для последующей работы остается пример Standalone кластера и Docker-образ для Kubernetes кластера. Следует так же отметить, что образ Spark может быть использован и для запуска контейнера с Spark Master. Кратко обо всех способах запуска Spark, как в локальном и standalone режимах, так с использованием Kubernetes, можно прочесть здесь, а о различиях менеджеров Yarn и Mesos можно прочесть здесь, или изучить вопрос самостоятельно. Необходимо загрузить архив со Spark с официального сайта. Ввиду проблем с совместимостью именно в моей программно-аппаратной конфигурации, мне пришлось использовать версию 3.2.1, хотя, на момент тестирования (и написания данной статьи) уже доступна версия 3.3.0. Распаковать содержимое в директорию spark (или воспользоваться подготовленными примерами из репозитория). После распаковки архива директория spark должна иметь следующий вид: За исключением директории rapids. Ее нужно создать и загрузить в нее файлы *.jar с сайта Rapids. На момент написания статьи доступен релиз 22.06.0, который совмещает в себе два представленных на скриншоте файла. Но на момент тестирования свежей версией была 22.04.0.  Сначала я хотел написать, что оставляю этот момент без изменений, однако, когда я тестировал Spring сервис перед публикацией статьи на 1080 Ti, я все же попробовал использовать 22.06.0. На 1080 Ti все равно не заработало, однако последняя версия вывела сообщение для дебага, благодаря которому я узнал, что новая версия Rapids в связке с Pascal и WSL2 работать не будет. Читатель может использовать любую версию из упомянутых, а в репозитории с примером остается 22.06.0. Скрипт getGpusResources.sh нужен для обнаружения GPU ресурсов: Еще одна директория - datasets. В ней хранятся файлы *.csv и *.parquet, которые в последующем будут использованы в приложениях как обучающие и валидирующие датасеты. Взять можно здесь. Пройдемся по всем директориям, в которых нужно внести изменения.  Все рабочие конфиги представлены в репозитории примера. Директория содержит шаблоны конфигов. Задействовать каждый можно путем копирования шаблона в ту же директорию и удаления “.template” в имени файла: Таким образом, редактируется файл spark-defaults.conf: spark-env.sh: Следуя документации Spark, следующим шагом должен быть запуск скрипта для создания Docker-образов: Подготовленный Dockerfile необходимо изменить до вида: entrypoint.sh: SPARK_DRIVER_BIND_ADDRESS - адрес Spark Master, в моем случае - это адрес локальной виртуальной машины. Для дебага возможных неисправностей в ходе работы с библиотекой nccl следует выставить уровень дебага INFO. Командой start-worker.sh spark://$SPARK_DRIVER_BIND_ADDRESS запускается воркер, подключаясь к мастеру. Исходники kubernetes/dockerfiles/Dockerfile и содержимое той же директории можно найти в репозитории примера. Сборка образа и запуск контейнера: Следует убедиться, что в образ установилась библиотека nccl корректной версии, для чего нужно экзекнуться в созданный контейнер Spark Worker и выполнить: на момент написания статьи и тестирования корректной версией является 2.12.12-1+cuda11.7. В версиях ниже может встречаться проблема с запуском задач XGBoost, т.к. nccl не может найти сетевое устройство по причине того, что в докер-контейнере оно является виртуальным.  Проверяем доступность воркера, путем открытия его WEB GUI по адресу localhost:8081 (в соответствии с командой docker run выше): Видим, что помимо Cores и Memory, доступен ресурс Resources: gpu. На моей локальной машине одно устройство, и его id обозначен в массиве как “0”. Проверяем WEB GUI мастера (адрес моей локальной виртуальной машины http://192.168.5.129:8080/): Spark Worker, запущенный в контейнере, появился в списке Workers. Можно переходить к приложению. В текущем разделе рассматривается пример работы простого веб-сервиса, являющимся так же Spark Driver. Приложение будет иметь 3 HTTP Endpoint’a, на каждом будет доступен пример одной из библиотек: DJL, Spark ML, XGBoost.  В виде каркаса приложения используется Spring Boot с зависимость spring-boot-starter-web, используемый JDK - OpenJDK 8 (держим в уме требование Rapids). Я создаю новый проект в тот момент, когда пишу эту статью, поэтому финальный результат также должен заработать у читателя при условии выполнения предварительных шагов, описанных выше. Структура проекта: Файл pom.xml можно посмотреть в репозитории, я остановлюсь на важном моменте. Для работы XGBoost на Windows с WSL2 в Docker контейнере проведено детальное обследование проблемы в GitHub Issue. На данный момент версия библиотеки XGBoost, используемая в данном примере, не имеет релизной версии, поэтому в Maven central она отсутствует. Для загрузки библиотеки нужно добавить в pom.xml репозиторий с версиями SNAPSHOT: Однако, есть нюанс. Доступа с Российских и Казахстанских IP (с других не проверялось) к данному репо с недавнего времени нет. Варианты: либо VPN, либо воспользоваться репозиторием проекта и загрузить джарники xgboost4j-gpu_2.12-2.0.0-SNAPSHOT.jar и xgboost4j-spark-gpu_2.12-2.0.0-SNAPSHOT.jar в локальный m2 репозиторий: Также эти джарники необходимо загрузить в директорию jars проекта (см. скрин выше). Данные *.jar файлы будут переданы в Spark Executor как зависимости для запуска кода драйвера. Список таких файлов описывается в SparkConfiguration:   Параметров конфигурации Spark очень много, подробнее с ними можно ознакомиться на странице Configuration - Spark 3.3.0 Documentation . Контроллер максимально простой, он содержит три сервиса, каждый сервис реализует по 1-2 метода каждой библиотеки. Обращаю внимание, что данный контроллер является инструментом запуска соответствующего примера, сделанный в угоду скорости и самому факту, что несколько технологий можно объединить в приложении Spring, и никак не претендует на использование приложения в продуктивной среде. Для реального приложения здесь должны быть как минимум другие HTTP глаголы, обработчики сообщений, информативные DTO, асинхронные операции, брокеры сообщений для потоков данных, реактивщины, вебсокеты и вот это вот все. Первая библиотека на очереди - DJL. Это удобная библиотека машинного обучения для языка Java, особенностью которой является зоопарк моделей (Model Zoo), позволяющий получить готовую модель по описываемым параметрам из списка доступных моделей. Также имеется возможность создать свою модель, сохранить на диск и загрузить для дальнейшего использования. В данном примере рассматривается реализация модели линейной регрессии. К сожалению, в виду архитектурной особенности данной модели, распараллелить процесс ее обучения довольно сложно, и, вероятно, решается в определенных случаях определенными движками, такими как PyTorch. По крайней мере, распараллеливание обучения с помощью Spark модели линейной регрессии мне не попалось, и быстро сам придумать реализацию также не смог. Однако, есть распространенный пример применения Spark в паре с DJL для классификации изображений с использованием модели из Model Zoo, например данная статья. Реализация модели линейной регрессии сделана на основе статей 3.2. Linear Regression Implementation from Scratch — Dive into Deep Learning 0.1.0 documentation  и 3.3. Concise Implementation of Linear Regression — Dive into Deep Learning 0.1.0 documentation и отображена в сервисе DJLService. В целях отладки в main методе приложения логируется вызов нескольких методов, с помощью которых легко опознать некорректность конфигурации приложения. При корректной конфигурации должен вывестись лог вида: Код примера:   Модель DataPoints: Результат выполнения: Существует замечательная документация для начинающих от nVidia по работе со Spark ML на примере модели Random Forest. Учитывая специфику данной модели, процесс обучения можно распараллелить на несколько исполнителей, а затем пользоваться либо средним значением в случае решения задач регрессии, либо голосованием по большинству в случае решения задач классификации. Подробнее можно почитать на хабре, в документации Spark, примеры кода Spark ML также можно посмотреть в документации. В данном примере понадобятся датасеты для тренировки и валидации, можно взять отсюда, либо воспользоваться кодом репозитория. Отмечу, в данном разделе не полностью переписан пример из статьи nVidia по Spark ML, а скорее является реализацией задачи из статьи nVidia по XGBoost, но с применением Random Forest из Spark ML. Датасеты копируются в сценарии Dockerfile, а в сервисе пути к ним хардкодятся (пример же, можно себе позволить). Обращаю внимание: в статье по Spark ML говорится, что только XGBoost поддерживает GPU-ускорение в Spark ML. Вполне может быть, что документация устарела (как писали в одном из Issue на GitHub) и в данный момент, так как в документации Rapids указывается репозиторий с еще как минимум одним примером для алгоритма Principal component analysis (PCA). Код сервиса: Последним примером является реализация примера из статьи nVidia по XGBoost, который использует и Spark, и Rapids вместе. Данный пример является самым интересным, так как обеспечивает действительно лучшую скорость вычислений по сравнению со Spark ML Random Forest.  Кроме того, в документации Rapids первым примером рассматривается операция Join двух датафреймов из 10 млн чисел. Данный пример также реализован в тестовом методе сервиса RapidsService: Пример несколько отличается от своего исходника на Scala, но также обеспечивает вычисления на GPU. DAG представлен на скрине ниже: Что касается XGBoost, то пример взят из статьи nVidia, датасеты те же, что и для Random Forest Spark ML, про сам XGBoost можно почитать здесь и здесь. Реализация XGBoost regressor: Пора запустить приложение. Для этого необходимо собрать Docker-образ, воспользовавшись Dockerfile: Обращаю внимание на первую и вторую строки. Если нужно запускать приложение с логикой библиотеки DJL, нужно воспользоваться базовым образом, созданным ранее для Spark. Он содержит необходимое системное ПО для работы с видеокартами nVidia. Учитывая, что при запуске скачивается указанный в зависимостях движок (PyTorch, MXNet, etc), нужно соединение с интернетом и немного больше времени на запуск сервиса. Есть вариант один раз подключить volume к контейнеру и добавить несколько директив COPY в Dockerfile после первого запуска, чтобы не загружать необходимые файлы из интернета, а сразу скопировать в образ. В случае, если DJL не используется в сервисе, смысла в использовании “тяжелого” базового образа, и можно воспользоваться более “легким” образом, который содержит JRE8, например, тот, что закомментирован в первой строке. Команды запуска вынесены в скрипт build.sh: Через некоторое время контейнер должен запуститься, движки и все необходимое для DJL инициализироваться, а в Web UI Spark Master должен появиться сервис в списке запущенных приложений: В Web UI Spark Worker должен появиться Executor для указанного приложения:   Web UI сервиса также должен стать доступным:   Согласно контроллеру, имеется 4 доступных GET-метода, которые запускают необходимый пример: http://localhost:9090/djl http://localhost:9090/forest http://localhost:9090/gpu_test http://localhost:9090/xgboost Лог результата выполнения DJL представлен в соответствующем разделе выше, он не очень интересен, нежели Spark Jobs. При первом запуске логики, которая должна выполняться на Spark, метод выполняется несколько дольше, чем при последующих запусках - есть необходимость некоторого “прогрева” executor'а. Для этого можно запустить метод gpu_test. В деталях джобы видим, что она выполнялась чуть более 8и секунд, что, действительно, довольно долго. Последующие вызовы данного метода выполнялись в два раза быстрее (кроме второго - именно в данный момент подвел vmmem, выделив под WSL 25Гб ОЗУ):   Результат выполнения:   Random forest сделал очень много Spark Jobs, которые выполнялись с 14:04:02 до 14:05:15 (73 секунды).  XGBoost на том же датасете выполнялся в рамках 433-436 Spark Jobs, которые заняли ~16 секунд. Результаты: Именно в данном примере видим, что XGBoost справился быстрее и лучше, судя по значениям ошибок. Разворачивая спойлер из первой части статьи, заработало действительно не все. На двух машинах с Docker Desktop завести целевую схему не удалось по причине невозможности синхронизации двух контейнеров разных машин друг с другом. Network=host не дает нужного результата, роуты и nginx proxy тоже, также настраивал iptables в контейнерах - безуспешно. Проблему можно решить воспользовавшись Docker Swarm, но все дело в том, что для корректной работы кластера все равно нужна хотя бы одна машина с ОС Linux, выступающей мастером. Естественно, я попробовал сделать схему с запуском мастера на виртуалке, прописывал роуты и направлял трафик со второй физической ноды на определенный порт первой, а на первой ноде прописывал роут с данного порта на виртуалку, но столкнулся с проблемой получения ответных пакетов от мастера, и несколькими другими проблемами. Также можно было попробовать раскатать Kubernetes, но на этом я решил остановиться, так как: а) Standalone кластер Spark в контейнерах - по сути бред и априори overhead, так как суть Standalone кластера заключается в том, что его можно использовать на малом количестве нод и для постоянной нагрузки. В таком случае Docker не нужен, и лучше поставить на чистую ОС; б) Если Kubernetes, то нужно понимать, что он нужен для плавающих нагрузок, для оптимизации использования вычислительных ресурсов, и лучше использовать Kubernetes Operator - вот в этом опыта пока еще нет, и, вероятно, это тема будущей статьи; в) “Все, стоп, осталось только кубер на винде раскатать, хватит страдать фигней” - раздалось в голове, и я остановился  (но в следующей следующей статье продолжил) :)  Однако, результат меня все равно обрадовал - работа XGBoost встала на шаге синхронизации датасетов для последующей выдачи результата, что успел запечатлеть на скриншоте. Цель данной статьи считаю достигнутой. Все три библиотеки оказались работоспособными, сервис написан на Java, запущен как Spring Web Service, в Docker-контейнерах задачи на GPU исполняются. Что дальше и можно ли что-то улучшить? Естественно, направлений работы несколько: Тюнинг Spark. Как минимум, неплохо бы подключить Kryo serializer. Во время работы с Rapids 22.06.0 у меня он так и не заработал. Кроме Kryo есть множество параметров конфигурации самого Spark, которые все вместе в целом довольно сильно влияют на производительность. Запуск Spark Standalone кластера на bare metal и нативном Ubuntu 20.04. Запуск сервиса в Kubernetes в паре с Spark Kubernetes Operator. Вероятно, гайд по запуску и результаты будут темой отдельной статьи. Дальнейшее R&D в ML и Spark. Accelerating Spark 3.0 and XGBoost End-to-End Training and Hyperparameter Tuning Accelerating Deep Learning on the JVM with Apache Spark and NVIDIA GPUs How Amazon retail systems run machine learning predictions with Apache Spark using Deep Java Library How Netflix uses Deep Java Library (DJL) for distributed deep learning inference in real-time Adopting machine learning in your microservices with DJL (Deep Java Library) and Spring Boot Getting Started with RAPIDS Accelerator with on premise cluster or local mode Accelerating Apache Spark 3.0 with GPUs and RAPIDS Leverage deep learning in Scala with GPU on Spark 3.0 Accelerating Deep Learning on the JVM with Apache Spark and NVIDIA GPUs nVidia docker containers documentation CUDA on WSL User Guide How to install CUDA Toolkit on Ubuntu 18.04 LTS — Performatune WSL 2 GPU Support for Docker Desktop on NVIDIA GPUs — Docker nVidia Docker images Get Started — RAPIDS Docs On-Prem On-Prem — Example Join Operation Predictive Analytics Tutorial with Spark ML | NVIDIA What’s New in Deep Learning & Artificial Intelligence from NVIDIA Spark ML library documentation: Classification and regression — Spark 3.3.0 Documentation Ensembles — RDD-based API — Spark 3.3.0 Documentation Main — Deep Java Library Examples Troubleshooting — Deep Java Library Deep Learning with Spark in Deep Java Library in 10 minutes Deep Java Library(DJL) — a Deep Learning Toolkit for Java Developers 5.5. GPUs — Dive into Deep Learning 0.1.0 documentation DJL dependency management — Deep Java Library 3.2. Linear Regression Implementation from Scratch — Dive into Deep Learning 0.1.0 documentation 3.3. Concise Implementation of Linear Regression GitHub — dmlc/xgboost: Scalable, Portable and Distributed Gradient Boosting (GBDT, GBRT or GBM) Library, for Python, R, Java, Scala, C++ and more. Runs on single machine, Hadoop, Spark, Dask, Flink and DataFlow xgboost/jvm-packages/xgboost4j-example at master · dmlc/xgboost xgboost/SparkMLlibPipeline.scala at master · dmlc/xgboost A Full Integration of XGBoost and Apache Spark XGBoost4J-Spark-GPU Tutorial (version 1.6.1+) — xgboost 2.0.0-dev documentation XGBoost4J-Spark-GPU Tutorial (version 1.6.1+) — xgboost 1.6.1 documentation xgboost4j_2.12 1.6.1 API spark-rapids-examples/kubernetes-scala.md at branch-22.06 · NVIDIA/spark-rapids-examples spark-rapids-examples/Taxi.scala at branch-22.06 · NVIDIA/spark-rapids-examples How to extract best parameters from a CrossValidatorModel Use shared library that uses glibc on AlpineLinux