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

Введение

MLOps-экосистема ClearML состоит из множества компонентов и покрывает полный цикл разработки и доставки ML-моделей. Ранее мы разобрали работу с ClearML Data и ClearML Session, а в новой статье расскажем про 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. В него сохраняются:

• ссылка на Git-репозиторий;
• незакоммиченные изменения в виде diff-патча;
• автоматически собранный список зависимостей;
• параметры, метрики, stdout / stderr;
• конфигурации и артефакты.

Task — это полное описание состояния проекта на момент запуска.

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

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

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

Работа с ClearML начинается с регистрации на сервере и установки ClearML на локальную машину. После этого необходимо настроить ClearML Agent.

Эти шаги уже были проделаны в статье о ClearML Session. Чтобы не дублировать информацию, добавим ссылки на инструкции оттуда:

1. Инструкция по регистрации и установке ClearML;
2. Инструкция по настройке ClearML Agent.

Для статьи подготовлен репозиторий, в котором дообучается 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. В нём:

1. Устанавливается clearml-agent.
2. Проверяется доступность GPU.
3. Добавляется ssh-ключ, чтобы агент ClearML смог склонировать github-репозиторий с обучением.

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

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

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

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

Далее нужно:

1. Авторизоваться на github.com, перейти в настройки: Settings → SSH and GPG keys → New SSH key.
2. Скопировать содержимое ключа с расширением .pub и вставить в поле «Key». После этого нажать кнопку «Add SSH key»:
   

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

2. В Google Colab подключается Google Drive, при этом 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.

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

Пример создания отдельных 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, клонировав ранее выполненный эксперимент;
• из кода через SDK;
• из терминала, используя CLI.

Назначение задачи в очередь средствами 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. Агент, прослушивающий очередь, извлекает задачу и запускает её выполнение с нуля.

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

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

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

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

• —base-task-id — id задачи, на основе которой будет воспроизведён эксперимент;
• —queue — очередь, на которую назначаем задачу. Важно, чтобы был агент, который её слушает.

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

Логи выполнения запущенной любым из перечисленных способов задачи будут доступны в ClearML UI.

Заключение

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

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

Полезные ссылки

1. Документация ClearML
2. Репозиторий с проектом из статьи
3. ClearML в Google Colab
4. ClearML Data Management
5. ClearML Session
6. Пример использования execute_remotely
7. Репозиторий ClearML agent
8. Полезные примеры использования ClearML Task CLI