ClearML Agent: обучение модели в Google Colab

ClearML — это целый космос, так что мы продолжаем разбирать его компоненты. В прошлой статье мы рассматривали ClearML Session и настраивали удаленную среду разработки с Jupyter Lab и VSCode. В этот раз поговорим о ClearML Agent и разберем, как с его помощью запустить обучение на удаленном сервере, в частности, Google Colab.

Итак, поехали!

Кратко о ClearML Agent

ClearML Agent — это виртуальная среда и менеджер выполнения решений DL/ML на машинах с GPU «в одном флаконе». Он представляет собой полноценное кластерное решение.

Вот основные возможности ClearML Agent:

• запуск экспериментов на произвольных ресурсах — локальных или в облаке;

• воспроизведение экспериментов: агент клонирует код, устанавливает зависимости, воспроизводит окружение, а затем запускает код (это позволяет запускать задачи на удаленной машине так же, как и на локальной);

• запуск сервисов/контейнеров для долгоживущих задач (не только коротких экспериментов);

• управление ресурсами: агент можно запустить в CPU-only режиме или с --gpus, а также управлять выделением GPU через NVIDIA_VISIBLE_DEVICES;

• планирование выполнения задач: можно запускать задачи по очереди/в зависимости от очередей, а также назначать им приоритеты;

• масштабирование: можно добавлять/удалять машины из «кластера», использовать ресурсы разных машин — облачных и on-premise вместе;

• логирование: stdout/stderr процесса, установка зависимостей, клонирование — всё логируется и может просматриваться через UI;

• конфигурация запуска/переопределение параметров задачи через UI, SDK, CLI — без необходимости менять код, например, можно менять гиперпараметры, зависимости, аргументы командной строки, docker-образ;

• возможность интеграции с Kubernetes/bare-metal/SLURM.

Если вкратце, ClearML Agent обеспечивает запуск кода на удаленной машине, воссоздавая окружение проекта так, словно эксперимент проводится локально. Заглянем под капот этого процесса.

1. При локальном запуске кода вызов Task.init(...) создает новую Task на ClearML Server. В этой Task сохраняется:

   • ссылка на Git-репозиторий;

   • незакоммиченные изменения в виде diff-патча;

   • автоматически собранный список зависимостей;

   •  параметры, метрики, stdout / stderr;

   •  конфигурации и артефакты.

   Поясним, что Task — это полное и исчерпывающее описание состояния проекта в момент запуска.

2. После локального запуска Task отправляется в очередь — канал, по которому задачи распределяются от сервера к агентам. Процесс назначения может осуществляться через SDK, CLI или ClearML UI, о чем расскажем подробнее ниже. Агент способен слушать сразу несколько очередей, откуда задачи берутся в работу. Каждый агент работает только над одним task-ом в моменте, сколько бы их не скопилось в очереди. Притом ClearML Agent выполняет следующие действия.

• Загружает метаданные задачи с ClearML Server, получая план действий от главного вычислительного центра.

• Клонирует Git-репозиторий, указанный в задаче, воссоздавая основу проекта. Если репозиторий приватный, необходимо предоставить SSH-ключи или credentials.

• Применяет diff-патч с локальными незакомиченными изменениями, настраивая репозиторий под конкретную задачу. Патч применяется только к текстовым файлам.

• Создает виртуальное окружение (pip / conda / poetry), подготавливая лабораторное пространство для эксперимента.

• Восстанавливает зависимости (сначала из списка библиотек, собранного на локальной машине, затем из requirements.txt, если он существует).

• Запускает основной скрипт в этом окружении, оживляя код.

• Передает stdout, метрики и логи обратно на ClearML Server.

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

Подготовка окружения

Самое время перейти от теории к практике. Работа с ClearML начинается с регистрации на сервере и установки экосистемы на локальную машину. Далее необходимо настроить ClearML Agent и создать очереди. Мы уже рассматривали эти два шага в статье про ClearML Session, так что здесь повторяться не станем и сразу пойдем дальше.

Поскольку в одной статье все не отразишь, мы подготовили несколько полезных «плюшек» для подготовки окружения. Прежде всего это репозиторий с дообучением BERT на решение определения тональности текста. В README.md содержится подробный гайд по локальному развертыванию проекта.

Но и это не все: в файле train.py реализовано дообучение модели с использованием фреймворка PyTorch Lightning и логированием в ClearML.

Локальный запуск обучения и логирование в ClearML

Опишем пошагово логирование эксперимента в ClearML.

# Импортируем Task из библиотеки clearml
from clearml import Task

BATCH_SIZE = args.batch_size
MAX_EPOCHS = args.max_epochs
LEARNING_RATE = args.learning_rate
MAX_LENGTH = 128
SUBSET_SIZE = args.subset_size

# Самая главная часть для логгирования Task.init 
# с Task.init уже работает логгирование большей части информации в ClearML Server
task = Task.init( # 
    project_name="BERT-Sentiment-Classification",
    task_name="IMDB-Training-BERT-Base",
    task_type=Task.TaskTypes.training,
)

# логгируем интересующие нас гиперпараметры
task.connect(
      {
          "batch_size": BATCH_SIZE,
          "max_epochs": MAX_EPOCHS,
          "learning_rate": LEARNING_RATE,
          "max_length": MAX_LENGTH,
          "subset_size": SUBSET_SIZE,
          "model_name": "bert-base-uncased",
          "device_arg": args.device,
          "cuda_available": device_info["cuda_available"],
          "device_name": device_info["device_name"],
          "device_count": device_info["device_count"],
      }
  )

# Так как мы используем PyTorch Lightning, добавляем логгер clearml в класс модели
model = BERTSentimentClassifier(
    learning_rate=LEARNING_RATE,
    num_labels=2,
    clearml_logger=task.get_logger(),  # передаем ClearML logger
)

Как говорится, семь раз отмерь: прежде чем запускать код на удаленном сервере, сделаем это локально. Установим зависимости согласно инструкции в README.md и запустим файл train.py для начала обучения:

python train.py

Далее проверяем в ClearML UI корректность логирования.

Запуск ClearML Agent в Google Colab

Для начала создадим очередь, откуда в агент будут поступать задачи. Это можно сделать в ClearML UI в разделе с очередями «Queues». Если не хочется совершать лишних телодвижений, используйте канал по умолчанию — default.

Далее создаем агента, который будет слушать эту очередь. Для выполнения его действий подготовлен ноутбук в Google Colab, в котором:

• устанавливается clearml-agent;

• проверяется доступность GPU;

• добавляется SSH-ключ, чтобы агент ClearML смог склонировать GitHub-репозиторий с обучением.

Инструкция по созданию SSH-ключа и добавлению его на GitHub

Генерация SSH-ключей выполняется следующей командой:

ssh-keygen -t ed25519 -C "example@gmail.com"

В результате будут сгенерированы два файла: id_ed25519.pub и id_ed25519.

Далее нужно:

• Авторизоваться на GitHub, перейти в настройки: Settings → SSH and GPG keys → New SSH key.

• Скопировать содержимое ключа с расширением .pub и вставить в поле «Key». После этого нажать кнопку «Add SSH key»:

Теперь пройдем весь путь создания агента. 

Шаг 1. Добавляем созданные SSH-ключи в папку на Google Drive.

Шаг 2. Подключаем Google Drive в Google Colab. При этом Google запрашивает подтверждение для доступа к вашим данным:

from google.colab import drive
drive.mount('/content/drive')

Шаг 3. Подключаем и проверяем доступ к GitHub:

# копируем эти файлы из google drive в папку .ssh в Google Colab
!mkdir /root/.ssh/
!cp -r /content/drive/MyDrive/rsa_keys/id_ed25519* /root/.ssh/

# Проверяем 
!ls -a /root/.ssh

# Для получения публичных ключей хоста SSH для GitHub.com добавляем github в файл known_hosts
!ssh-keyscan github.com >> /root/.ssh/known_hosts

# Проверяем, что мы можем получить доступ к нашему github
!ssh -T git@github.com # Hi {тут будет ваш никнейм}! You've successfully authenticated, but GitHub does not provide shell access.

# Добавляем переменную окружения, чтобы не отрисовывать и не выводить информацию matplotlib
!export MPLBACKEND=TkAg 

Шаг 4. Устанавливаем credentials для ClearML Agent.

Как и на этапе настройки ClearML, хорошей практикой считается создание отдельных credentials для агента.

Пример создания отдельных credentials для ClearML Agent

Settings → Workspace → Create new credentials

Добавляем переменные CLEARML_API_ACCESS_KEY и CLEARML_API_SECRET_KEY в секретах Google Colab:

Затем эти переменные используются для подключения к ClearML Server.

from clearml import Task
from google.colab import userdata

# Читаем переменные окружения
clearml_api_key = userdata.get('CLEARML_API_ACCESS_KEY')
clearml_secret_key = userdata.get('CLEARML_API_SECRET_KEY')

# Настраиваем креды для подключения к ClearML Server
Task.set_credentials(
    api_host="<https://api.clear.ml>",
    web_host="<https://app.clear.ml>",
    files_host="<https://files.clear.ml>",
    key=clearml_api_key,
    secret=clearml_secret_key,
)

Шаг 5. Запускаем ClearML Agent для прослушивания очереди. По умолчанию задача выполняется в фоновом режиме.

Список аргументов для запуска агента доступен по ссылке.

Поведение агента задается переменными окружения, которые отображаются в терминале при запуске. Задать или поменять значение переменной можно в файле clearml.conf. Или же это делается напрямую — через переменную окружения. Список переменных доступен по ссылке. Будьте внимательны: не все переменные в clearml.conf дублируются переменными окружения!

Итак, ClearML Agent на удаленном сервере создан и отслеживается, о чем сообщается в ClearML UI. Теперь на агента можно назначать задачи.

Назначение задачи в очередь

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

Назначение задачи в очередь средствами ClearML UI

Появившуюся в ClearML UI задачу можно клонировать, после чего она перейдет в статус Draft. Это позволяет менять гиперпараметры задачи и назначать её на выполнение в нужную очередь.

На видео ниже продемонстрировано, как это сделать:

Назначение задачи в очередь с помощью SDK

В этом случае для добавления задачи в очередь необходимо дописать в код train.py следующую строку:

task.execute_remotely(queue_name="deep_school")

Вуаля! При выполнении скрипта train.py создается Task в ClearML Server, где сохраняется всё необходимое для воспроизведения эксперимента. Вызов task.execute_remotely() завершает локальное выполнение кода и назначает созданный task из ClearML Server в очередь deep_school, указанную в аргументе queue_name. Прослушивающий очередь агент извлекает задачу и запускает её выполнение с нуля.

Лайфхак: лучше вызвать task.execute_remotely() после первой или второй эпохи, когда нет сомнений в том, что эксперимент работает локально и не падает.

Назначение задачи в очередь с помощью CLI

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

clearml-task --base-task-id b8e466bc55b24a5b9cc164b6cb986fab --queue deep_school

После выполнения этой команды агент возьмет задачу из очереди и запустит её в полностью воспроизводимой среде:

• —base-task-id — id задачи, на основе которой будет воспроизведен эксперимент;

• —queue — очередь, на которую назначаем задачу. Важно, чтобы был агент, который её слушает.

Документация по ClearML Task доступна по ссылке.

Вне зависимости от способа запуска логи выполнения задачи доступны в интерфейсе ClearML UI.

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

Среди плюсов ClearML Agent — возможность запускать эксперименты на разных удаленных серверах c логированием в ClearML UI. Важно, что при запуске задачи в ClearML применяются все изменения в репозитории (даже незакоммиченные). Это позволяет сравнивать эксперименты и дает уверенность в их воспроизводимости.

P.S. Первая версия статьи написана для блога DeepSchool. На Хабре публикуется с изменениями и дополнениями.