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

• запуск экспериментов на произвольных ресурсах — локальных или в облаке;
  
• воспроизведение экспериментов: агент клонирует код, устанавливает зависимости, воспроизводит окружение, а затем запускает код (это позволяет запускать задачи на удаленной машине так же, как и на локальной);
  
• запуск сервисов/контейнеров для долгоживущих задач (не только коротких экспериментов);
  
• управление ресурсами: агент можно запустить в CPU-only режиме или с --gpus, а также управлять выделением GPU через NVIDIA_VISIBLE_DEVICES;
  
• планирование выполнения задач: можно запускать задачи по очереди/в зависимости от очередей, а также назначать им приоритеты;
  
• масштабирование: можно добавлять/удалять машины из «кластера», использовать ресурсы разных машин — облачных и on-premise вместе;
  
• логирование: stdout/stderr процесса, установка зависимостей, клонирование — всё логируется и может просматриваться через UI;
  
• конфигурация запуска/переопределение параметров задачи через UI, SDK, CLI — без необходимости менять код, например, можно менять гиперпараметры, зависимости, аргументы командной строки, docker-образ;
  
• возможность интеграции с Kubernetes/bare-metal/SLURM.

• При локальном запуске кода вызов Task.init(...)создает новую Task на ClearML Server. В этой Task сохраняется:
  • ссылка на Git-репозиторий;
    
  • незакоммиченные изменения в виде diff-патча;
    
  • автоматически собранный список зависимостей;
    
  •  параметры, метрики, stdout / stderr;
    
  •  конфигурации и артефакты.
  Поясним, что Task — это полное и исчерпывающее описание состояния проекта в момент запуска.
  
• После локального запуска Task отправляется в очередь — канал, по которому задачи распределяются от сервера к агентам. Процесс назначения может осуществляться через SDK, CLI или ClearML UI, о чем расскажем подробнее ниже. Агент способен слушать сразу несколько очередей, откуда задачи берутся в работу. Каждый агент работает только над одним task-ом в моменте, сколько бы их не скопилось в очереди. Притом ClearML Agent выполняет следующие действия.

• Загружает метаданные задачи с ClearML Server, получая план действий от главного вычислительного центра.
  
• Клонирует Git-репозиторий, указанный в задаче, воссоздавая основу проекта. Если репозиторий приватный, необходимо предоставить SSH-ключи или credentials.
  
• Применяет diff-патч с локальными незакомиченными изменениями, настраивая репозиторий под конкретную задачу. Патч применяется только к текстовым файлам.
  
• Создает виртуальное окружение (pip / conda / poetry), подготавливая лабораторное пространство для эксперимента.
  
• Восстанавливает зависимости (сначала из списка библиотек, собранного на локальной машине, затем из requirements.txt, если он существует).
  
• Запускает основной скрипт в этом окружении, оживляя код.
  
• Передает stdout, метрики и логи обратно на ClearML Server.

# Импортируем 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
)

python train.py

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

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

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

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

Рисунок 5. Создание SSH-ключей
В результате будут сгенерированы два файла: id_ed25519.pub и id_ed25519.
Далее нужно:

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

Рисунок 6. Добавление публичного SSH-ключа в GitHub для доступа к приватным репозиториям поSSH

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

# копируем эти файлы из 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

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

Settings → Workspace→ Create new credentials

Рисунок 8. Настройка рабочего пространства и создание кредов

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,
)

task.execute_remotely(queue_name="deep_school")

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

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