【MLflow】Tailscale + AWS EC2/S3で作る、セキュアなLLM学習管理環境の構築ガイド
はじめに
LLM(大規模言語モデル)のファインチューニングを行う際、実験ごとのハイパーパラメータやLossの推移、モデルの保存(Artifacts)を管理するために MLflow は欠かせないツールです。
今回は、手元のオンプレミスGPUサーバー(DGX Sparkなど)またはクラウド上のGPUインスタンスで学習する場合のいずれでも、ログやモデルデータを一元的に管理する目的でAWS(EC2上のMLflow + S3)に集約するアーキテクチャを構築しました。
さらに、セキュリティを高めるために Tailscale を導入してVPN経由でアクセスするようにしたのですが、環境構築でいくつかハマりました。
この記事では、MLflowでML/LLMOps環境を作ろうとしている方に向けて、アクセス権限・Dockerビルド・MLflowのHost Header拒否問題 などの躓きポイントと、その最終的な解決策を共有します。
全体アーキテクチャ
- 学習環境(クライアント): オンプレミスのGPUサーバー(DGX Sparkなど)
- トラッキングサーバー: AWS EC2 (Ubuntu) + Docker Compose
- バックエンドDB(メタデータ): PostgreSQL (Dockerコンテナ内)
- アーティファクト保存先: AWS S3バケット
- ネットワーク: Tailscale (VPN) を用いてセキュアに通信
[オンプレGPU] --(Tailscale VPN)-- [AWS EC2 (MLflow Server)] --(IAM Role)--> [AWS S3]
躓きポイント1:S3へのアクセス権限(IAMロールのベストプラクティス)
S3へモデル(Artifact)を保存するためにはAWSの認証が必要です。最初は .env にアクセスキーを書こうかと考えましたが、セキュリティの観点から EC2のIAMロール を使用するのが正解です。
Boto3(AWS SDK)が勝手にロールの権限を拾ってくれるため、コードや .env にクレデンシャルを書く必要がなくなります。
{
"Version": "2012-10-17",
"Statement": [
{
"Sid": "ListObjectsInBucket",
"Effect": "Allow",
"Action": ["s3:ListBucket"],
"Resource": ["arn:aws:s3:::your-mlflow-bucket"]
},
{
"Sid": "AllObjectActions",
"Effect": "Allow",
"Action": ["s3:GetObject", "s3:PutObject", "s3:DeleteObject"],
"Resource": ["arn:aws:s3:::your-mlflow-bucket/*"]
}
]
}
躓きポイント2:Dockerビルド時の can't stat pgdata エラー
Docker ComposeでPostgreSQLとMLflowを立ち上げる際、DBのデータをホストのディレクトリ(./pgdata)にバインドマウントしていました。
すると、次回ビルド時に error checking context: can't stat '/path/to/pgdata' というエラーでDockerがクラッシュしました。
原因: PostgreSQLコンテナが root 権限で pgdata を作成したため、ホスト側の一般ユーザーから読み取れず、Dockerのビルドコンテキストのスキャンがコケていた。
解決策: Dockerが管理する Named Volume(名前付きボリューム) に変更することで、スキャンの対象外となり、アクセス権限エラーも消滅しました。
(※古いDocker Compose v1を使っていると発生しやすい問題でもあり、Compose V2へのアップデートも必須です)
##最大の沼:Invalid Host header - possible DNS rebinding attack detected
TailscaleのIPアドレス(100.x.x.x)でブラウザからアクセスしようとすると、真っ白な画面にこのエラーが表示され、アクセスが弾かれました。
なぜ起きたのか?(MLflowのセキュリティ強化)
MLflow v2.11以降/v3系で導入されたセキュリティ強化により、以下の2つの条件が揃うと外部からのアクセスをDNSリバインディング攻撃とみなしてシャットアウトします。
- 本番用サーバー(Gunicorn)ではなく、開発用サーバー(Werkzeug) で起動してしまっている。
- ホストヘッダーのホワイトリスト(
--allowed-hosts)が設定されていない。
解決策(Dockerfile & entrypoint.sh)
YAMLの command に複雑な引数を書くとパースエラーになりやすいため、起動スクリプト(entrypoint.sh) に切り出し、強制的にGunicornで立ち上げる構成にしました。
FROM ghcr.io/mlflow/mlflow:v3.8.1
# S3アクセス用のboto3、DB用のpsycopg2、そして本番サーバー用のgunicornをインストール
RUN pip install --no-cache-dir psycopg2-binary boto3 gunicorn
#!/bin/bash
set -e
echo "Starting MLflow Server with Gunicorn..."
# --workers を指定してGunicornを強制
# --allowed-hosts "*" でTailscale経由のアクセスを許可
exec mlflow server \
--backend-store-uri $BACKEND_STORE_URI \
--artifacts-destination $ARTIFACT_ROOT \
--serve-artifacts \
--host 0.0.0.0 \
--port 5000 \
--workers 2 \
--allowed-hosts "*"
version: '3.8'
services:
db:
image: postgres:14
restart: always
environment:
POSTGRES_USER: ${POSTGRES_USER}
POSTGRES_PASSWORD: ${POSTGRES_PASSWORD}
POSTGRES_DB: ${POSTGRES_DB}
volumes:
- db_data:/var/lib/postgresql/data # Named Volumeを使用
mlflow:
build: .
image: mlflow-gunicorn-server:v1.0
restart: always
ports:
- "5000:5000"
environment:
AWS_DEFAULT_REGION: ap-northeast-1
BACKEND_STORE_URI: postgresql://${POSTGRES_USER}:${POSTGRES_PASSWORD}@db:5432/${POSTGRES_DB}
ARTIFACT_ROOT: ${ARTIFACT_ROOT}
volumes:
- ./entrypoint.sh:/entrypoint.sh
entrypoint: ["/bin/bash", "/entrypoint.sh"]
depends_on:
- db
volumes:
db_data:
これで、sudo docker compose up -d --build を実行すれば、堅牢で安定したMLflowサーバーがTailscale上に誕生します。
HF Transformersとの連携(LLMの学習ロギング)
サーバー側のセットアップさえ完了すれば、学習環境側の実装は非常にシンプルです。
Hugging Faceの Trainer を使ってLLMを学習させる場合、コードに数行追加するだけで自動的にMLflowへログが送信されます。
import mlflow
from transformers import Trainer, TrainingArguments, AutoModelForCausalLM, AutoTokenizer
# 1. 構築したTailscale上のMLflowサーバーを指定
mlflow.set_tracking_uri("http://mlflow-server:5000")
mlflow.set_experiment("LLM_Finetuning_Experiment")
# 2. 自動ロギングを有効化 (Lossや学習率、パラメータを自動記録)
# ※モデルサイズが大きい場合は log_models=False にしてアーティファクトのS3アップロードを防ぐ
mlflow.transformers.autolog(log_models=False)
# モデルとデータセットの準備
model = AutoModelForCausalLM.from_pretrained("...")
# ...
# 3. Trainerの引数に report_to="mlflow" を指定
args = TrainingArguments(
output_dir="./results",
num_train_epochs=3,
report_to="mlflow", # ★ココ
)
trainer = Trainer(
model=model,
args=args,
# ...
)
trainer.train()
これだけで、Lossカーブの描画、学習環境のハイパーパラメータ、使用したデータセット情報などがすべて一元管理されるようになります。
まとめ
- 認証: S3連携はアクセスキーを避け、EC2のIAMロールを使う。
- Docker構成: DBの永続化はNamed Volumeを使い、古いComposeを使わない。
-
MLflow起動:
Invalid Host header回避のため、gunicornの明示的インストールと--allowed-hosts "*"の指定が必須。 -
運用: 複雑な起動コマンドは
entrypoint.shにまとめることでデバッグが劇的に楽になる。
オンプレミスのGPUリソースと、クラウドGPUインスタンスのマネージドなスケーラビリティを組み合わせたハイブリッド環境は、ML/LLM開発において非常に強力です。こうしたハイブリッドGPU環境でのMLflow環境構築を考えている方の参考になれば幸いです!
Discussion