Привет! Меня зовут Александр, я MLOps-инженер.

JupyterHub — централизованный инструмент для создания Jupyter-ноутбуков для разных пользователей по заранее заданным параметрам, который используется более чем 200 специалистами у нас на проекте с коллегами из Альфа-Банка.

Сейчас множество сервисов Альфа-Банка могут запускаться в Kubernetes, который уже стал стандартом. И JupyterHub не исключение. Есть много статей по его запуску, и в большинстве из них описано, как просто поднять данный сервис и запустить рабочий ноутбук на базе профилей. Каждый профиль представляет собой заранее настроенные параметры для Pod ноутбука.

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

Оглавление

Проблема веса образов и ограничений в разработке
Проблема невозможности динамического выделения ресурсов под ноутбуки
Решение проблемы ноутбуков
Решение проблемы с динамическими ресурсами
Итог

Мы изучили много материалов о том, как поднимать Jupyter, и посмотрели на множество образов. Обычно все Docker-образы Jupyter Notebooks наследуемые, и сейчас есть два основных источника: Jupyter Docker Stacks и Kubeflow Notebooks. Образы имеют примерно одинаковую структуру:

Базовый образ
Базовый образ с Jupyter
Образы с различными библиотеками, разделённые на категории:
- Jupyter SciPy
- Jupyter PyTorch
- Jupyter TensorFlow

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

Проблема веса образов и ограничений в разработке

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

Более того, разные пользователи разрабатывают разные модели, и разные модели могут использовать разные версии библиотек. Но при использовании нашего кастомного образа со всеми библиотеками мы не только ставим пользователя в жёсткие рамки, но и поддерживаем очень тяжёлый образ с большим количеством библиотек определённых версий. Размеры данного образа, особенно если он предназначен для использования CUDA, могут достигать порядка 20 ГБ. Но даже и без использования драйверов для Nvidia вес мог достигать 15 ГБ, что по общепринятым меркам является заоблачным значением.

Таким образом, образов либо много, либо есть общий, но со слишком большим объёмом. К тому же мы фиксируем огромный набор библиотек, которые могут попросту не использоваться.

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

Как уже упоминалось ранее, в стандартном запуске Jupyter пользователи получают некоторые профили, в которых можно определить разные параметры для ноутбуков. Однако что делать, если эти параметры имеют большой разброс?

Например, одному пользователю может потребоваться 10 ГБ памяти, а другому — 500 ГБ. Такой же разброс может возникнуть и с CPU, и с GPU.

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

Вы разделили профили на 10 ГБ, 50 ГБ и 100 ГБ. Но разработчику может потребоваться, например, 60 ГБ. Ему придётся выбирать профиль, который резервирует на определённой ноде 100 ГБ, 40 ГБ из которых будут простаивать и не могут быть использованы, если только не завершить работу ноутбука. Если наши пользователи используют определённую ноду в кластере, то в худшем случае мы фактически теряем 40% оперативной памяти на данной ноде.

В дополнение к выделению ресурсов, множество разработчиков часто разделяются на различные группы, каждая из которых занимается определёнными задачами. Некоторым нужно использование GPU, в то время как другие совсем без него обходятся. Или, возможно, определённой группе требуется работа на выделенных узлах Kubernetes для максимальной загрузки ресурсов, в то время как другим это не нужно.

При обычном запуске JupyterHub не предоставляет никаких возможностей для работы с различными группами пользователей, которые, например, могут быть определены в LDAP.

Решение проблемы ноутбуков

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

Легковесность. Ранее, когда у нас были образы со всеми библиотеками, их размеры достигали порядка 15-20 ГБ, а их деление на разные группы умножало это число.
Неограниченность в выборе библиотек. Идея ноутбуков заключается в том, чтобы предоставить пользователям пространство для разработки, не ограничивая их в выборе версии Python и его библиотек. Решить, какие библиотеки нужны - задача DS. В то время как MLOps инженер предоставляет возможность установить эти библиотеки.
Интеграция с инфраструктурой. Образы должны содержать все инструменты, чтобы в рамках вашего контура, ваш разработчик мог спокойно скачать любую библиотеку или подключиться к какому-либо сервису, например, MLFlow и Minio, минимально разбираясь в инфраструктуре.

Требованиям легковесности удовлетворяет базовый образ с Jupyter, однако в него стоит добавить возможность работы с окружениями и интеграцию с инфраструктурой. Для интеграции с инфраструктурой в исходный образ мы включили различные дополнительные утилиты, системные пакеты, драйверы для NVidia, чтобы работали GPU на кластере, Java для Spark, before-start скрипты, переменные окружения до популярных сервисов вроде MLFlow и другое. Все это формируется исходя из ваших потребностей.

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

Здесь важным моментом является подключение Volumes при запуске. Нам необходимо подключить домашнюю директорию пользователя. Это необходимо, чтобы он не терял данные при перезапуске. Также необходимо настроить mamba так, чтобы установка новых окружений происходила также в домашний каталог.

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

Наша основная задача — предоставить пользователю готовый к работе образ, чтобы он поднял ноутбук, выполнил pip install -r requirements.txt для конкретной модели в отдельном окружении и сразу же приступил к работе.

Решение проблемы с динамическими ресурсами

Теперь давайте перейдём к решению проблем с ресурсами и использованием пользовательских групп.

JupyterHub имеет под собой множество компонентов. Среди них есть Spawner, и для Kubernetes он имеет название KubeSpawner. Этот компонент позволяет запускать ваши ноутбуки в виде Kubernetes Pods. По умолчанию он предоставляет пользователю UI форму с заранее заготовленными профилями.

А что если мы изменим эту форму и предоставим пользователям свою?

Давайте кратко рассмотрим, что происходит, когда вы авторизуетесь, выбираете профиль и создаёте ноутбук. Как только происходит авторизация, начинается работа спавнера, который предоставляет весь интерфейс. И как только пользователь выбирает какой-либо профиль из списка и нажимает на кнопку "Start", KubeSpawner считывает конфигурацию выбранного профиля и на её основе запускает Pod.

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

Давайте построим план действий.

Первое, что необходимо сделать — это создать класс для нашего нового Spawner, который будет унаследован от стандартного KubeSpawner.

# alexspawner/spawner/spawner.py
class CustomSpawner(KubeSpawner):

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

        self.logger = setup_logger(__name__, logging.ERROR)
        self.logger.info('Start working with CustomSpawner')

В конструкторе класса мы можем описать логику обработки групп, так как Spawner стартует сразу после авторизации.

# Вызов внутри CustomSpawner.__init__() ↴
# self.user_groups = get_user_groups(self.logger, self.user.name)

# alexspawner/spawner/utils.py
def get_user_groups(logger, username) -> list:
    """
    Узнаем какой группе принадлежит пользователь в LDAP.
    Либо получаем какие-либо другие данные, на основе которых будут 
      действовать разделения прав пользователей.
    """

Далее нам необходимо переписать метод _options_form_default, в котором идёт генерация формы. Форму мы создавали через Jinja2 Template, как это сделано и в оригинальном KubeSpawner.

# alexspawner/spawner/spawner.py
def _options_form_default(self):
    #  Реализация функции render_template и формы в репозитории в конце статьи
    #  alexspawner/spawner/utils.py
    form = render_template(self.groups_data_config, self.group_for_render)
    return form

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

# alexspawner/spawner/spawner.py
def options_from_form(self, formdata):
        image = formdata.get('jupyter_image', [''])[0].strip()
        self.image = select_image_from_input(image)

        cpu = formdata.get('cpu', [''])[0].strip()
        self.cpu_guarantee = round(float(cpu) / 3, 2)
        self.cpu_limit = float(cpu)

        mem = formdata.get('mem', [''])[0].strip()
        self.mem_guarantee = str(mem) + "G"
        self.mem_limit = str(mem) + "G"

        options = {
            'image': self.image,
            'cpu': self.cpu_limit,
            'mem': self.mem_limit,
        }

        return options

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

Значение Mem Limits = Mem Requests, в то время как CPU Requests — это лишь треть от указанных лимитов на форме. Сделано это по той причине, что нагрузка на ноутбуки идёт непостоянная. Из графиков потребления мы видим, что, когда пользователь пишет код, CPU в целом не нагружается, и скачки происходят только в момент, когда пользователь запускает различные процессы, такие как обучение.

Потребление CPU при работе в Jupyter Notebook с моделями машинного обучения

В дополнение ко всему, комбинируя решение проблемы с ноутбуками и новый кастомный Spawner, мы можем задать поле с выбором образов для запуска, таких как: stable, stable-old, hotfix, dev. На данный момент при любых изменениях в чарте будет перезагружаться Pod Hub (GitHub Issue), что является большой проблемой, так как убивает процессы в ноутбуках. С помощью унифицированных тегов мы сможем не только уменьшить количество перезагрузок, но и безболезненно проводить обновления наших легковесных образов Jupyter.

Итог

Теперь, когда у нас есть переработанный Spawner, пользователи могут более компактно размещаться на нодах Kubernetes. В соответствии с группой, одни имеют возможность использовать GPU, а другие — нет. Кто-то может заезжать только на "общие" ноды кластера, а кто-то потребляет максимальные возможные ресурсы для обучения своей модели уже на выделенных машинах. Для возможности отслеживания доступных ресурсов у нас создан Grafana Dashboard.

В следующем GitHub репозитории: dc24_custom_kubespawner вы можете ознакомиться с минимальной версией CustomSpawner для Kubernetes, который изменяет интерфейс пользователя JupyterHub и позволяет ему выбирать нужные ресурсы CPU и памяти.

Другие наши статьи про DevOps:

Сравнение Kubernetes с другими оркестраторами: Docker Swarm, Nomad, Openshift
Как деплоить приложение на Django в Kubernetes с нуля
Как изучать Kubernetes джуну — и зачем
Как построить эргономичный мониторинг
Поднимаем динамические окружения для stateless- и stateful-сервисов
Зачем мы сделали собственный контроллер для копирования секретов в Kubernetes
Как экономить ресурсы, затрачиваемые Kubernetes

JupyterHub на стероидах: реализация KubeFlow фич без масштабных интеграций