Александр Рыжков

Ментор Skillfactory, руководитель команды LightAutoML и 4х Kaggle Grandmaster

MLflow — это инструмент, который стал для меня (и многих команд, с которыми я работал) настоящим спасением от хаоса в ML-проектах.

Многие начинающие в ML наверняка сталкивались с ситуацией: вы пробуете разные модели, меняете параметры, запускаете обучение снова и снова… и через пару дней уже не можете вспомнить, какой именно набор параметров дал тот самый лучший результат. Или, что еще хуже, вы получили отличную модель на своем ноутбуке, а у коллеги на его машине она не воспроизводится. На помощь придет MLflow.

Что такое MLflow и зачем он нужен?

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

MLflow — это opensource-платформа для управления всем жизненным циклом машинного обучения. Инструмент помогает навести порядок в экспериментах, сделать их воспроизводимыми, упростить упаковку кода и развертывание моделей.

Ключевые задачи, которые решает MLflow:

Отслеживание экспериментов: запись параметров, метрик, версий кода и артефактов (например, обученных моделей) для каждого запуска. Больше никаких потерянных результатов!
Воспроизводимость: упаковка кода и его зависимостей так, чтобы любой мог запустить ваш эксперимент и получить те же результаты (при условии тех же данных, конечно).
Управление моделями: централизованное хранилище для ваших обученных моделей с версионированием и управлением стадиями (например, «разработка», «тестирование», «продакшн»).
Развертывание моделей: MLflow предлагает инструменты для упрощения развертывания моделей в различных средах.

Думайте об MLflow как об организованном лабораторном журнале, системе контроля версий для моделей и конвейере для их доставки — все в одном флаконе.

MLflow для трекинга экспериментов — основные компоненты

Это, пожалуй, самая часто используемая часть MLflow, особенно на начальных этапах. Компонент mlflow.tracking позволяет логировать все самое важное во время обучения модели. Давайте разберемся с основными понятиями:

Experiment (Эксперимент). Это как папка для связанных задач. Например, у вас может быть эксперимент «Прогнозирование оттока клиентов» или «Классификация изображений котов». Внутри одного эксперимента вы проводите множество запусков.
Run (Запуск). Один конкретный запуск вашего кода. Например, обучение модели с определенным набором гиперпараметров. Каждый раз, когда вы запускаете скрипт обучения с логированием в MLflow, создается новый Run.
Parameters (Параметры). Входные данные вашего запуска, которые вы хотите сохранить. Обычно это гиперпараметры модели (скорость обучения, глубина дерева), версии используемых данных, фичи.
Metrics (Метрики). Числовые значения, которые показывают качество вашего запуска. Это могут быть точность (accuracy), ошибка (loss), F1-score. Метрики могут меняться со временем (например, loss по эпохам).
Artifacts (Артефакты). Любые файлы, которые вы хотите сохранить как результат запуска. Самый частый артефакт — это сама обученная модель. Также это могут быть графики (например, learning curve), файлы с предсказаниями, примеры данных.

Все это добро удобно просматривать во встроенном веб-интерфейсе MLflow UI. Вы можете запустить его командой mlflow ui в терминале, и он покажет вам все ваши эксперименты, запуски, позволит сравнивать их, сортировать по метрикам и параметрам. Это невероятно удобно для анализа результатов!

Что такое MLflow Model Registry?

Итак, вы провели кучу экспериментов и нашли несколько отличных моделей (артефактов). Что дальше? Как управлять ими? Как понять, какая модель сейчас используется в продакшене, а какая только тестируется? Здесь на помощь приходит MLflow Model Registry (реестр моделей).

Model Registry позволяет:

Централизованно хранить ваши лучшие модели.
Версионировать модели. Если вы обучили новую, улучшенную версию той же модели, она получит новый номер версии (например, ChurnPredictor v2).
Присваивать моделям стадии (Stages):

Staging: Модель на тестировании, кандидат на переход в продакшен.
Production: Модель, которая используется в реальной системе.
Archived: Устаревшая модель, больше не используется.
Можно добавлять и свои стадии.

Добавлять описания и теги к моделям и их версиям.
Отслеживать происхождение модели (из какого Run она была создана).

Это делает процесс перехода модели от разработки к продакшену гораздо более контролируемым и прозрачным. Команда DevOps или другие разработчики могут легко понять, какую версию модели им нужно взять для развертывания.

Проекты в MLflow — как использовать? Создание конфига проекта

Компонент MLflow Projects — это стандартный формат для упаковки вашего кода машинного обучения так, чтобы его можно было легко запустить где угодно — на вашей машине, на машине коллеги или в облаке.

Основная идея — описать ваш проект в специальном файле MLproject (в формате YAML). Этот файл содержит:

name: Имя вашего проекта (необязательно).
environment: Указание, как создать окружение для запуска кода. Чаще всего используется ссылка на файл conda.yaml, который описывает все необходимые Python-библиотеки и их версии. Также можно использовать Docker-контейнер.
entry_points: Точки входа — команды, которые можно выполнить в рамках проекта. Обычно есть основная точка входа (main), которая запускает ваш скрипт обучения или предсказания. Здесь же можно определить параметры, которые можно передать в скрипт при запуске.

Пример простого MLproject-файла:

name: tutorial

environment: conda.yaml

entry_points:
  main:
    parameters:
      alpha: {type: float, default: 0.5}
      l1_ratio: {type: float, default: 0.1}
    command: "python train.py --alpha {alpha} --l1_ratio {l1_ratio}"

А вот пример conda.yaml:

name: tutorial-env
channels:
  - defaults
  - conda-forge
dependencies:
  - python=3.9
  - pip
  - pip:
    - mlflow>=1.0
    - scikit-learn
    - pandas

Теперь, имея эти файлы и ваш скрипт (train.py), любой человек (или система CI/CD) может запустить ваш проект командой:

mlflow run . --experiment-name="tutorial" -P alpha=0.4 -P l1_ratio=0.2

MLflow автоматически создаст указанное окружение (если его еще нет), загрузит зависимости и выполнит команду из entry_points, передав указанные параметры.

Практический пример использования MLflow на Python с проектом и Model Registry

Создадим простой проект, который обучает линейную регрессию на случайных данных, логирует все в MLflow Tracking, использует MLflow Projects для запуска и регистрирует модель в Model Registry.

Шаг 1: Структура проекта

Создайте папку, например mlflow_example, и внутри нее следующие файлы:

mlflow_example/
 ├── train.py
 ├── MLproject
 └── conda.yaml

Шаг 2: Файл conda.yaml (окружение)

name: mlflow-sklearn-example
 channels:
   - defaults
   - conda-forge
 dependencies:
   - python=3.9
   - pip
   - pip:
     - mlflow>=1.10 # Используем версию с поддержкой Model Registry
     - scikit-learn
     - numpy
     - pandas

Шаг 3: Файл MLproject (конфиг проекта)

name: sklearn-regression-example
 
 environment: conda.yaml
 
 entry_points:
   main:
     parameters:
       n_samples: {type: int, default: 100}
       noise: {type: float, default: 0.1}
       model_name: {type: str, default: "LinearRegressionDemo"} # Имя для регистрации модели
     command: "python train.py --n_samples {n_samples} --noise {noise} --model_name {model_name}"

Шаг 4: Файл train.py (скрипт обучения)

import argparse
import numpy as np
from sklearn.linear_model import LinearRegression
from sklearn.model_selection import train_test_split
from sklearn.metrics import mean_squared_error

import mlflow
import mlflow.sklearn

def generate_data(n_samples, noise):
    X = np.random.rand(n_samples, 1) * 10
    y = 2 * X.squeeze() + 3 + np.random.randn(n_samples) * noise * 10
    return train_test_split(X, y, test_size=0.2, random_state=42)

def train_model(n_samples: int, noise: float, model_name: str):
    """Обучает модель и логирует параметры, метрику и артефакт."""

    # Генерируем данные
    X_train, X_test, y_train, y_test = generate_data(n_samples, noise)

    # Логируем параметры
    with mlflow.start_run(nested=True): 
        mlflow.log_param("n_samples", n_samples)
        mlflow.log_param("noise", noise)
        mlflow.log_param("model_name", model_name)

        # Обучаем
        model = LinearRegression()
        model.fit(X_train, y_train)

        # Предсказываем и логируем MSE
        y_pred = model.predict(X_test)
        mse = mean_squared_error(y_test, y_pred)
        mlflow.log_metric("mse", mse)
        print(f"Mean Squared Error: {mse}")

        # Логируем модель как артефакт и регистрируем её
        mlflow.sklearn.log_model(
            sk_model=model,
            artifact_path="sklearn-model",
            registered_model_name=model_name
        )
        print(f"Model logged and registered under name: {model_name}")

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--n_samples", type=int, default=100)
    parser.add_argument("--noise", type=float, default=0.1)
    parser.add_argument("--model_name", type=str, default="LinearRegressionDemo")
    args = parser.parse_args()

    train_model(
        n_samples=args.n_samples,
        noise=args.noise,
        model_name=args.model_name
    )

if __name__ == "__main__":
    main()

Шаг 5: Запуск и просмотр результатов

Установите MLflow: pip install mlflow (если еще не установлен).
Перейдите в папку mlflow_example в вашем терминале.
Запустите MLflow UI в фоновом режиме (или в другом окне терминала):

mlflow ui
Откройте в браузере http://127.0.0.1:5000.

Запустите проект:

mlflow run . --experiment-name="sklearn-regression-example" -P n_samples=200 -P noise=0.2 -P model_name="LinRegNoise0.2"
MLflow создаст conda-окружение (это может занять время при первом запуске), а затем выполнит train.py с указанными параметрами.

Проверьте MLflow UI:

Вы увидите новый запуск (Run) в эксперименте sklearn-regression-example.
Внутри запуска будут залогированы параметры (n_samples, noise, model_name_param).
Будет залогирована метрика mse.
В разделе Artifacts вы увидите папку sklearn-model, содержащую файлы модели.
Перейдите в раздел Models в MLflow UI. Вы увидите зарегистрированную модель LinRegNoise0.2 с версией 1.

Попробуйте запустить проект еще раз с другими параметрами, например:

mlflow run . --experiment-name="sklearn-regression-example" -P n_samples=2000 -P noise=0.4 -P model_name="LinRegNoise2000_0.4"

Вы увидите новый Run и новую зарегистрированную модель. Если вы используете то же model_name, MLflow создаст новую версию существующей модели.

Трекинг LLM c помощью MLflow

С взрывным ростом больших языковых моделей (LLM) возникает вопрос: а можно ли использовать MLflow для них? Ответ — да! Принципы остаются теми же, но акценты смещаются.

При работе с LLM с помощью MLflow можно отслеживать:

Параметры:

Имя базовой модели (e.g., gpt-3.5-turbo, llama-2-7b).
Параметры генерации (temperature, top_p, max_tokens).
Версию промпта (или сам текст промпта, если он не слишком длинный).
Параметры fine-tuning (если вы дообучаете модель): learning rate, batch size, количество эпох.

Метрики:

Стандартные метрики оценки текста (BLEU, ROUGE).
Метрики оценки на конкретных задачах (точность классификации, F1 для извлечения сущностей).
Оценки качества ответа, полученные с помощью других LLM (например, оценка релевантности ответа).
Latency (время ответа модели).
Стоимость вызова API (если используется платная модель).

Артефакты:

Полный текст промптов (особенно если они сложные или шаблонные).
Примеры сгенерированных ответов (input/output pairs).
Наборы данных для fine-tuning или оценки.
Визуализации (если применимо).
Логи ошибок или неудачных генераций.
Если вы делаете fine-tuning, то саму дообученную модель (например, адаптеры LoRA).

MLflow также развивает интеграции с популярными LLM-фреймворками (например, LangChain, Transformers), чтобы сделать логирование еще проще. Основная идея — использовать знакомые mlflow.log_param, mlflow.log_metric, mlflow.log_artifact для специфичных для LLM сущностей.

Обучиться работе с моделями машинного обучения: от базовой математики до написания собственного алгоритма — можно на совместной магистратуре Skillfactory и МИФИ «Прикладной анализ данных и машинное обучение».

Что такое MLFlow и как он помогает в разработке моделей