Запуск ML скриптов в облаке с помощью dstack. Бонус – про запуск open-source проектов

Привет! Пару недель назад мы выложили на GitHub утилиту для запуска ML скриптов в облаке. 

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

---

Итак, по порядку

Команда dstack состоит из трех разработчиков. До этого , я работал в JetBrains, в командах IntelliJ IDEA, DataGrip, GoLand и PyCharm. Можно глянуть другие мои посты на Хабре, где я рассказывал про работу в JetBrains.

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

Немного о проблеме 

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

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

ssh -i /path/key-pair-name.pem instance-user-name@instance-public-dns-name

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

О модных MLOps платформах

Те, кто не хочет изобретать собственные велосипеды (для автоматизации запуска скриптов), часто обращают внимание на MLOps платформы (AWS SageMaker, Vertex AI, и др.).

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

В классической разработке есть Git, Docker, Terraform и т.п.. В машинном обучении настолько универсальных инструментов довольно мало (исключение – DVC).

Если вы предпочитаете работать с Python скриптами, использовать Git, инструменты и библиотеки от разных вендоров, и вообще работать из своей любимой IDE, вы можете найти MLOps платформы скорее контрпродуктивными. 

Какую задачу решает dstack

Если запускать ML скрипты в облаке, на что тратится время:

1. Создание и настройка удаленной машины

2. Копирование кода и данных

3. Настройка окружающей среды

Дополнительное неудобство связано с необходимостью использовать SSH в качестве основного интерфейса.

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

Представьте, если бы вы могли запускать ML скрипты в облаке как если бы вы их запускали локально. А создание, настройка, и остановка удаленных машин делалась бы автоматически. В таком случае разработка могла бы вестись в вашей любимой IDE без необходимости использовать какие либо MLOps платформы.

Эту задачу и решает утилита dstack.

Вкратце, как работает dstack

1. Вы описываете требования ML скриптов к ресурсам и окружающей среде в .dstack/workflows.yaml, внутри вашего проекта

2. Вы описываете в ~/.dstack/config.yaml в каком S3 bucket хранить метаданные и результаты работы скриптов

3. Вы запускаете сприпты с помощью команды dstack run, и мониторите их выполнение с помощью команды dstack ps.

Вот небольшой пример файла .dstack/workflows.yaml:

workflows: 
  - name: "train"
    provider: bash
    deps:
      - tag: mnist_data
    commands:
      - pip install requirements.txt
      - python src/train.py
    artifacts: 
      - path: checkpoint
    resources:
      interruptible: true
      gpu: 1

Если запустить скрипт командой dstack run, команда сама создаст нужные ресурсы в AWS, настроит окружение, скачает исходники и данные, запустит скрипт, сохранит результаты и затем удалит ресурсы в AWS.

Для вас это будет вылядеть, как будто вы запускаете скрипт src/train.py локально.

dstack run train

Provisioning... It may take up to a minute. ✓

To interrupt, press Ctrl+C.

Epoch 4: 100%|██████████████| 1876/1876 [00:17<00:00, 107.85it/s, loss=0.0944, v_num=0, val_loss=0.108, val_acc=0.968]
`Trainer.fit` stopped: `max_epochs=5` reached. 
Testing DataLoader 0: 100%|██████████████| 313/313 [00:00<00:00, 589.34it/s]

Test metric   DataLoader 0
val_acc       0.965399980545044
val_loss      0.10975822806358337

Команда dstack ps позволяет мониторить запущенные скрипты (по аналогии с Docker):

dstack ps -a

RUN               TARGET    STATUS  APPS  ARTIFACTS   SUBMITTED    TAG
angry-elephant-1  download  Done          data        8 hours ago  mnist_data
wet-insect-1      train     Running       checkpoint  1 weeks ago

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

В скриптах можно не только обрабатывать данные и тренировать модели, но и запускать приложения (FastAPI, Streamlit, Gradio, и т.п.) и среды разработки (например VS Code или JupyterLab)

Если кому-то интересно попробовать dstack, добро пожаловать в документацию!

Про трудности разработки 

Где хранить стейт

Одна из проблем, с которыми мы столкнулись, это хранение стейта утилиты (какие скрипты были запущены, в каком они состоянии, и т.п.) Традиционных подход заключается в использовании базы данных и в выделении центрального фасада (backend) для работы с этим стейтом. 

Если использовать такой подход, то прежде чем пользователь сможет использовать утилиту, он должен был бы самостоятельно установить и настроить такой фасад (а если проще, сервер). 

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

Скорость запуска в облаке

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

От момента выполнения команды dstack run до старта выполнения скрипта должно быть минимум времени.

Для этого нам потребовалось настроить пайплайн сборки собственного образа облачной машины (известного также как AMI). Для сборки образа мы воспользовались утилитой Packer, которая вызывается прямо из GitHub Actions в момент сборки dstack.

Отображение результатов выполнения

Еще одой интересной проблемой, с которой мы столкнулись, оказалось отображение результатов выполнения скрипта. Дело в том, что ML скрипты часто используют библиотеки типа tqdm – для отображения прогресса выполнения задач.

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

Это заставило нас попотеть. Зато в результате у пользователя создается полное ощущение, что скрипт выполняется локально.

Про запуск

Здесь я бы хотел поделиться опытом запуска проектов в целом и open-source проектов в частности.

Перестать откладывать

Самым неочевидным в запуске оказался вопрос “когда релизить – сейчас или немного позже”. Вот несколько причин, из-за которых мы постоянно откладывали релиз, и из-за чего потеряли много времени:

1. Хочется доделать “вот эту” и “вот ту” фичи и только потом релизить

2. Хочется сделать “красивый” сайт прежде чем релизить

3. Хочется пофиксить все баги и только потом релизить

4. Хочется релизить в “правильный” день (например, четверг)

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

Про коммьюнити

До недавнего времени у меня не было четкого пониманию, какую конкретно роль играет комьюнити, и как это использовать для запуска open-source продукта.

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

Про буст

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

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

Про отношение

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

Если интересно поговорить про dstack, запуск ML скриптов, или про разработку open-source проектов, пишите в личное сообщения! Буду рад поделиться опытом, или помочь, если нужна помощь.