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

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

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

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

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

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

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

1. Отслеживание экспериментов: запись параметров, метрик, версий кода и артефактов (например, обученных моделей) для каждого запуска. Больше никаких потерянных результатов!
2. Воспроизводимость: упаковка кода и его зависимостей так, чтобы любой мог запустить ваш эксперимент и получить те же результаты (при условии тех же данных, конечно).
3. Управление моделями: централизованное хранилище для ваших обученных моделей с версионированием и управлением стадиями (например, «разработка», «тестирование», «продакшн»).
4. Развертывание моделей: 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 позволяет:

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

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

4. Добавлять описания и теги к моделям и их версиям.
5. Отслеживать происхождение модели (из какого 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: Запуск и просмотр результатов

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

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

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

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

5. Проверьте 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 сущностей.

Коротко об MLFlow

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

Даже если вы только начинаете, я настоятельно рекомендую попробовать интегрировать MLflow в ваши учебные проекты. Вы сразу почувствуете разницу, когда через неделю сможете легко найти тот самый удачный запуск и понять, чем он отличался от остальных. По мере роста сложности ваших задач возможности MLflow Projects и Model Registry станут незаменимыми помощниками.