Отслеживание запусков с помощью MLflow и страницы запуска заданий

Это важно

Эта функция доступна в общедоступной предварительной версии.

Каждая рабочая нагрузка, которую вы отправляете с помощью air run, является одновременно запуском задания Databricks и запуском MLflow:

  • Выполнение задания (видимое на странице заданий и конвейеров рабочей области) отслеживает выполнение: состояние, вычисления, повторные попытки и выходные данные драйвера.
  • Запуск MLflow отслеживает эксперимент: параметры, метрики, системные метрики и артефакты.

Одна отправка создает одно выполнение задания и один запуск MLflow. Повторная попытка создает новый запуск MLflow.

Эксперименты и запуски

Два поля YAML задания управляют тем, как запуск отображается в MLflow:

experiment_name: my-training # Creates or appends to this MLflow experiment
mlflow_run_name: baseline-lr3e5 # Names the MLflow run for this submission
compute:
  num_accelerators: 8
  accelerator_type: GPU_8xH100
command: torchrun --nproc_per_node=8 train.py
max_retries: 2
  • experiment_name (Обязательно): создает эксперимент MLflow с таким именем, если он не существует, или добавляет новый запуск к существующему эксперименту. Эксперимент включает множество запусков.
  • mlflow_run_name (Необязательно): задает имя запуска. Если не указано, по умолчанию используется имя эксперимента в качестве имени запуска (experiment_name).
  • max_retries (Необязательно): каждая попытка повтора — это новый запуск MLflow в одном эксперименте, поэтому можно сравнить попытки. Исходная отправка и её повторные попытки относятся к одному запуску задания.

Страница запуска MLflow с метриками

Вы можете перейти к запуску из трех мест:

  • Jobs: На странице запусков заданий отображается список ваших запусков, и для каждого запуска приводятся ссылки на соответствующие запуск и эксперимент MLflow.
  • MLflow: на странице "Эксперименты" перечислены эксперименты MLflow.
  • Предыдущие запуски: air get run <job-run-id> выводит кликабельные ссылки на задание, эксперимент и запуск MLflow. air list runs выводит список предыдущих запусков и позволяет фильтровать для поиска определенного запуска.
air get run <job-run-id> # Links to the job, experiment, and MLflow run
air list runs # List previous runs; filter to find a specific run

Системные метрики

Метрики системы gpu, ЦП и памяти записываются автоматически для каждого запуска. Конфигурация не требуется. Просмотрите их на вкладке системных метрик запуска MLflow.

Вкладка системных метрик запуска MLflow (GPU/ЦП/память)

Регистрировать пользовательские метрики

Платформа создает запуск MLflow и предоставляет его идентификатор процессу обучения с помощью переменной MLFLOW_RUN_ID среды. Используйте API для отслеживания MLflow, чтобы записывать собственные параметры, метрики и артефакты в этом запуске.

В распределенных рабочих нагрузках (с несколькими узлами) каждый узел использует один и тот же запуск MLflow. Записывайте в журнал только из процесса rank-0, чтобы каждая метрика записывалась один раз:

import os

import mlflow

# Log from rank 0 only; all nodes share the same MLFLOW_RUN_ID.
if os.environ.get("RANK", "0") == "0":
    with mlflow.start_run(run_id=os.environ["MLFLOW_RUN_ID"]):
        mlflow.log_param("learning_rate", 3e-4)
        for step, loss in enumerate(training_losses):
            mlflow.log_metric("train_loss", loss, step=step)

Журналы и артефакты

Просматривайте в потоковом режиме или скачивайте журналы запуска с помощью air logs:

air logs <job-run-id> # Stream logs from node 0
air logs <job-run-id> --node 2 # Logs from a specific node
air logs <job-run-id> --download-to ./logs/ # Download instead of streaming

Журналы также доступны в виде артефактов в рамках запуска MLflow. Чтобы сохранить контрольные точки модели, запишите их в том каталога Unity. О шаблонах создания контрольных точек и управлении томами см. в разделе «Отслеживание экспериментов и наблюдаемость».

Дополнительные ресурсы