Posted inAI

DockerでMLflowをデプロイする：実際のMLOpsパイプラインでML実験を追跡し、モデルレジストリを管理する

AI tutorial - IT technology blog
AI tutorial - IT technology blog

深夜2時の悪夢：あのモデル、どれだったっけ？

午前2時17分、メッセージが届いた——最新のモデルデプロイ後、本番環境の精度が大幅に低下したという。まず頭に浮かんだのはロールバックだったが、もっと深刻な問題があった：どのモデルバージョンが稼働しているのか、どのハイパーパラメータで学習させたのか、どのデータセット分割で作られたのか、まったく把握できていなかったのだ。2週間にわたって3台のマシンで数十回の実験を繰り返し、ファイル名はmodel_final.pkl、model_final_v2.pkl、そして定番のmodel_ACTUALLY_final.pklといったものばかりだった。

あの夜は、Jupyterノートブック、gitログ、CSVファイルを掘り起こして5分で済むはずの情報を復元するのに4時間を費やした。機械学習プロジェクトをまともな実験追跡なしに本番へ出したことがあるなら、この感覚はよくわかるだろう。あのインシデントが、私がMLflowをきちんとセットアップするきっかけになった——実際の経験から言えば、本番MLシステムに関わる前にマスターしておくべき必須スキルの一つだ。

Dockerに触れる前に：MLflowが実際に何をするのか

MLflowは、MLのエンドツーエンドのライフサイクルを管理するオープンソースプラットフォームだ。魔法ではない——モデルと実験のための整理された記録管理システムだ。プラットフォームは4つのコンポーネントで構成される：

  • トラッキングサーバー——各トレーニング実行のパラメータ、メトリクス、アーティファクト（モデルファイル、グラフ）を記録する
  • モデルレジストリ——モデルのバージョン管理、ステージ管理、昇格（Staging → Production → Archived）を行う一元的なストア
  • MLflow Models——scikit-learn、PyTorch、TensorFlowなどに対応した標準パッケージングフォーマット
  • Projects——トレーニングコードの再現可能なパッケージング（あまり使われないが、チームには便利）

ほとんどのチームにとって、トラッキングサーバーとモデルレジストリだけで混乱の90%は解消される。今日デプロイするのはまさにこの2つだ——Docker Composeを使うことでセットアップが再現可能になり、誰かのラップトップ上で朽ちることもない。

Docker ComposeでMLflowをセットアップする

MLflowのバックエンドにPostgreSQLデータベース（コンテナ再起動後も実行メタデータを保持）と、S3互換アーティファクトストアとしてMinIO（モデルファイルをコンテナのファイルシステム内に置かない）を使って構築する。

プロジェクト構成

mlflow-docker/
├── docker-compose.yml
├── .env
└── mlflow/
    └── Dockerfile

.envファイル

# .env
POSTGRES_USER=mlflow
POSTGRES_PASSWORD=mlflow_secret
POSTGRES_DB=mlflow
MINIO_ROOT_USER=minioadmin
MINIO_ROOT_PASSWORD=minioadmin
MINIO_BUCKET=mlflow-artifacts
MLFLOW_S3_ENDPOINT_URL=http://minio:9000
AWS_ACCESS_KEY_ID=minioadmin
AWS_SECRET_ACCESS_KEY=minioadmin

docker-compose.yml

version: "3.8"

services:
  postgres:
    image: postgres:15
    environment:
      POSTGRES_USER: ${POSTGRES_USER}
      POSTGRES_PASSWORD: ${POSTGRES_PASSWORD}
      POSTGRES_DB: ${POSTGRES_DB}
    volumes:
      - postgres_data:/var/lib/postgresql/data
    healthcheck:
      test: ["CMD", "pg_isready", "-U", "${POSTGRES_USER}"]
      interval: 10s
      retries: 5

  minio:
    image: minio/minio:latest
    command: server /data --console-address ":9001"
    environment:
      MINIO_ROOT_USER: ${MINIO_ROOT_USER}
      MINIO_ROOT_PASSWORD: ${MINIO_ROOT_PASSWORD}
    ports:
      - "9000:9000"
      - "9001:9001"
    volumes:
      - minio_data:/data

  minio-init:
    image: minio/mc:latest
    depends_on:
      - minio
    entrypoint: >
      /bin/sh -c "
      sleep 5;
      mc alias set local http://minio:9000 ${MINIO_ROOT_USER} ${MINIO_ROOT_PASSWORD};
      mc mb local/${MINIO_BUCKET} --ignore-existing;
      exit 0;
      "
    environment:
      MINIO_ROOT_USER: ${MINIO_ROOT_USER}
      MINIO_ROOT_PASSWORD: ${MINIO_ROOT_PASSWORD}
      MINIO_BUCKET: ${MINIO_BUCKET}

  mlflow:
    image: ghcr.io/mlflow/mlflow:v2.13.0
    ports:
      - "5000:5000"
    environment:
      MLFLOW_S3_ENDPOINT_URL: ${MLFLOW_S3_ENDPOINT_URL}
      AWS_ACCESS_KEY_ID: ${AWS_ACCESS_KEY_ID}
      AWS_SECRET_ACCESS_KEY: ${AWS_SECRET_ACCESS_KEY}
    command: >
      mlflow server
      --backend-store-uri postgresql://${POSTGRES_USER}:${POSTGRES_PASSWORD}@postgres/${POSTGRES_DB}
      --default-artifact-root s3://${MINIO_BUCKET}/
      --host 0.0.0.0
      --port 5000
    depends_on:
      postgres:
        condition: service_healthy
      minio-init:
        condition: service_completed_successfully

volumes:
  postgres_data:
  minio_data:

スタックを起動する

docker compose up -d

# 全サービスが正常に起動していることを確認する
docker compose ps

# MLflow UI → http://localhost:5000
# MinIO コンソール → http://localhost:9001

初回起動時はPostgreSQLの初期化とMinIOのバケット作成に約30秒かかる。mlflowコンテナがhealthyと表示されたら、ポート5000でUIにアクセスできる。

最初の実験をログに記録する

ローカルマシン（またはトレーニングコードを実行する環境）に、MLflowクライアントとS3アーティファクトアップロード用のboto3をインストールする：

pip install mlflow boto3

基本的な実験追跡

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

os.environ["MLFLOW_S3_ENDPOINT_URL"] = "http://localhost:9000"
os.environ["AWS_ACCESS_KEY_ID"] = "minioadmin"
os.environ["AWS_SECRET_ACCESS_KEY"] = "minioadmin"

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

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)

for n_estimators in [50, 100, 200]:
    with mlflow.start_run(run_name=f"rf-{n_estimators}-trees"):
        model = RandomForestClassifier(n_estimators=n_estimators, random_state=42)
        model.fit(X_train, y_train)

        accuracy = accuracy_score(y_test, model.predict(X_test))

        # パラメータとメトリクスを記録する
        mlflow.log_param("n_estimators", n_estimators)
        mlflow.log_metric("accuracy", accuracy)

        # モデル自体をアーティファクトとして記録する
        mlflow.sklearn.log_model(
            sk_model=model,
            artifact_path="model",
            registered_model_name="iris-random-forest"
        )

        print(f"n_estimators={n_estimators} → accuracy={accuracy:.4f}")

このスクリプトを実行して、http://localhost:5000を開く。iris-classification実験の下に3つの実行が表示され、それぞれのパラメータ、メトリクス、そしてMinIOに保存されたモデルアーティファクトを確認できる。任意の実行をクリックすると、データセット情報、コードバージョン、環境など、再現に必要なすべての情報を含む完全なリネージュが確認できる。

モデルレジストリを管理する

上記の3つの実行は、registered_model_nameを通じてモデルを自動的に登録済みだ。次は、ステージを通じてバージョンを管理しよう——これがロールバックを4時間の悪夢ではなく30秒の作業にする鍵だ。

ステージを通じてモデルを昇格させる

from mlflow.tracking import MlflowClient

client = MlflowClient(tracking_uri="http://localhost:5000")

# モデルの全バージョンを一覧表示する
versions = client.search_model_versions("name='iris-random-forest'")
for v in versions:
    print(f"Version {v.version} — run_id={v.run_id} — stage={v.current_stage}")

# 最高性能のバージョンをStagingに昇格させる
client.transition_model_version_stage(
    name="iris-random-forest",
    version="3",  # 200本の木のバージョン
    stage="Staging"
)

# 検証が通過したら、Productionに移行する
client.transition_model_version_stage(
    name="iris-random-forest",
    version="3",
    stage="Production"
)

# 古い本番モデルをアーカイブする
client.transition_model_version_stage(
    name="iris-random-forest",
    version="1",
    stage="Archived"
)

推論用に本番モデルをロードする

import mlflow.sklearn
import os

os.environ["MLFLOW_S3_ENDPOINT_URL"] = "http://localhost:9000"
os.environ["AWS_ACCESS_KEY_ID"] = "minioadmin"
os.environ["AWS_SECRET_ACCESS_KEY"] = "minioadmin"

mlflow.set_tracking_uri("http://localhost:5000")

# 現在Productionタグが付いているモデルを常にロードする
model = mlflow.sklearn.load_model("models:/iris-random-forest/Production")

sample = [[5.1, 3.5, 1.4, 0.2]]
print(model.predict(sample))  # [0]

ここで重要なのはmodels:/iris-random-forest/Productionという指定だ——推論コードにバージョン番号をハードコードしない。バージョン4をProductionに昇格させると、次の予測呼び出しから自動的にそのモデルが使われる。ロールバックも同様にシンプルだ：古いバージョンをProductionに戻すだけでよい。

これで手に入ったもの

この時点で、完全なセルフホスト型MLOpsバックボーンが揃った：実験メタデータはPostgreSQLに（再起動後も保持、クエリ可能）、モデルアーティファクトはMinIOに（S3互換、スケーラブル）、そしてデプロイパイプラインに直接対応した昇格ステージを持つクリーンなモデルレジストリだ。すべてのトレーニング実行が監査可能——誰が、どのコードで、どのデータ分割で、どのメトリクスで学習させたかが一目でわかる。

次に深夜2時に何かが壊れても、MLflow UIを開いてモデル名でフィルタリングし、どのバージョンがProductionにあるか確認し、そのトレーニングメトリクスを前のバージョンと比較して、2分以内にロールバックできる。復旧できるインシデントと何時間も燃え続ける大火事の差はそこにある。

チームにとっての自然な次のステップは、CI/CDパイプラインをトレーニング成功後にモデルを自動登録するよう組み込み、メトリクスのしきい値をクリアしないとProductionへの昇格をブロックする仕組みを作ることだ——しかし、ここで構築した基盤はそれをきれいに処理できる。インフラはすでに整っている。ワークフローポリシーを定義するのはあなた自身だ。

Share:
Tags: