MLflow: как вести учёт экспериментов и хранить модели на своём сервере

Типичная картина без учёта экспериментов: поменял параметры в скрипте, запустил, результат посмотрел в терминале — и он потерялся. Через неделю уже не вспомнить, при каких настройках модель была лучше, и её приходится обучать заново. Так работать можно, пока запусков пять. Когда их пятьдесят — уже нет.

MLflow решает эту проблему: записывает каждый эксперимент (параметры, метрики, файлы), хранит версии моделей и показывает веб-интерфейс, где результаты можно сравнить в одном окне. Его можно поднять на своём сервере — тогда данные и модели никуда не уходят.

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

Для начала в MLflow важны две части:

Компонент	Что делает
Tracking	Записывает параметры, метрики и файлы (артефакты) каждого запуска обучения
Model Registry	Хранит готовые модели по версиям и помечает, какая из них сейчас рабочая

Есть и другие (например, Projects — упаковка кода обучения, чтобы другой человек или CI воспроизвёл результат), но без них можно обойтись.

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

Установка и первый запуск

Для учёбы я поднимал MLflow на своём сервере в Docker Compose вместе с PostgreSQL, где хранится история запусков.

Самый быстрый способ попробовать — без базы, всё в локальных файлах:

pip install mlflow
mlflow server --host 127.0.0.1 --port 5000

Интерфейс откроется на http://localhost:5000.

Для постоянной работы лучше Docker Compose с PostgreSQL. Учтите: в официальном образе MLflow нет драйвера PostgreSQL, его нужно доставить. Проще всего — свой небольшой образ:

FROM ghcr.io/mlflow/mlflow:v3.17.2
RUN pip install --no-cache-dir psycopg2-binary

И compose.yaml рядом с ним:

services:
  postgres:
    image: postgres:16
    environment:
      POSTGRES_DB: mlflow
      POSTGRES_USER: mlflow
      POSTGRES_PASSWORD: ${POSTGRES_PASSWORD}
    volumes:
      - postgres-data:/var/lib/postgresql/data

  mlflow:
    build: .
    command: >
      mlflow server
      --backend-store-uri postgresql://mlflow:${POSTGRES_PASSWORD}@postgres/mlflow
      --artifacts-destination /artifacts
      --host 0.0.0.0 --port 5000
    ports:
      - "127.0.0.1:5000:5000"
    volumes:
      - mlflow-artifacts:/artifacts
    depends_on:
      - postgres

volumes:
  postgres-data:
  mlflow-artifacts:

Пароль лежит в файле .env рядом (POSTGRES_PASSWORD=...), а не в самом compose-файле. Порт открыт только на 127.0.0.1: у MLflow по умолчанию нет входа по паролю, поэтому наружу его выставляют только через обратный прокси с авторизацией.

Версию образа лучше указывать конкретную, а не latest: тогда обновление не случится неожиданно.

Логирование экспериментов в коде

В скрипте обучения:

import mlflow
import mlflow.sklearn
from sklearn.datasets import load_iris
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import accuracy_score
from sklearn.model_selection import train_test_split

mlflow.set_tracking_uri("http://localhost:5000")
mlflow.set_experiment("iris-classifier")

params = {"n_estimators": 100, "max_depth": 10}

with mlflow.start_run(run_name="rf-exp-1"):
    mlflow.log_params(params)

    X, y = load_iris(return_X_y=True)
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

    model = RandomForestClassifier(**params, random_state=42)
    model.fit(X_train, y_train)

    accuracy = accuracy_score(y_test, model.predict(X_test))
    mlflow.log_metric("accuracy", accuracy)

    mlflow.sklearn.log_model(model, name="model")

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

Запускаете скрипт — в интерфейсе появляется строка с n_estimators, max_depth и accuracy. Запускаете с другими параметрами — добавляется ещё одна, их можно сравнивать.

Сравнение экспериментов в интерфейсе

Интерфейс MLflow показывает все запуски эксперимента и позволяет сравнить:

• параметры — какие значения были;
• метрики — точность, loss и т. п., в том числе графики по ходу обучения;
• артефакты — сохранённые файлы: модель, графики, примеры данных;
• версию кода — если запуск сделан из Git-репозитория, MLflow запоминает коммит.

Запуски можно фильтровать: например, оставить только те, где metrics.accuracy > 0.95, и посмотреть, что общего в их параметрах.

Model Registry: управление версиями

Когда модель получилась хорошей, её регистрируют в реестре:

mlflow.register_model(
    model_uri="runs:/<run_id>/model",
    name="iris-classifier",
)

<run_id> — номер запуска, он виден в интерфейсе. Каждая регистрация под тем же именем создаёт новую версию: 1, 2, 3…

Чтобы отметить рабочую версию, в современных версиях MLflow используют псевдонимы (aliases), например champion для боевой и challenger для кандидата. Старые стадии Staging/Production устарели.

from mlflow import MlflowClient

MlflowClient().set_registered_model_alias("iris-classifier", "champion", 2)

Приложение загружает модель по псевдониму — models:/iris-classifier@champion — и при смене версии код менять не нужно.

Типичные ошибки

Не записали параметр. Поменяли настройку, а в MLflow её не передали — потом не понять, чем запуски отличались. Правило: всё, что влияет на результат, записывать как параметр. Для популярных библиотек есть mlflow.autolog() — он сам записывает параметры и метрики.

Тяжёлые артефакты в базе. Метаданные (параметры, метрики) живут в PostgreSQL, а файлы моделей — в отдельном хранилище артефактов: на диске или в S3-совместимом хранилище. Для больших моделей сразу настраивайте S3, а не локальную папку.

Нет резервных копий. Если пропадёт база или папка артефактов, пропадёт вся история экспериментов. Бэкапьте и PostgreSQL (pg_dump), и хранилище артефактов.

Дальше: от эксперимента к продакшену

MLflow — первый шаг. Дальше обычно добавляют:

1. Сигнатуру модели — описание входных и выходных данных. MLflow сохраняет её вместе с моделью и проверяет входные данные при вызове.
2. Проверку качества перед выкладкой — новая версия получает псевдоним champion только если её метрика не хуже текущей.
3. Автоматический деплой — пайплайн CI/CD или GitOps забирает модель из реестра и выкатывает её в Kubernetes.

Свой сервер даёт главное — контроль: данные и модели остаются у вас.

Нужна помощь с приватной ML-моделью на своём сервере или в облаке? Напишите — обсудим архитектуру и оценим сложность.