👏

【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 にクレデンシャルを書く必要がなくなります。

iamポリシー設定例
{
    "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リバインディング攻撃とみなしてシャットアウトします。

  1. 本番用サーバー(Gunicorn)ではなく、開発用サーバー(Werkzeug) で起動してしまっている。
  2. ホストヘッダーのホワイトリスト(--allowed-hosts)が設定されていない。

解決策(Dockerfile & entrypoint.sh)

YAMLの command に複雑な引数を書くとパースエラーになりやすいため、起動スクリプト(entrypoint.sh) に切り出し、強制的にGunicornで立ち上げる構成にしました。

dockerfile
FROM ghcr.io/mlflow/mlflow:v3.8.1

# S3アクセス用のboto3、DB用のpsycopg2、そして本番サーバー用のgunicornをインストール
RUN pip install --no-cache-dir psycopg2-binary boto3 gunicorn

entrypoint.sh
#!/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 "*"

docker-compose.yml
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へログが送信されます。

train.py
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環境構築を考えている方の参考になれば幸いです!

株式会社Quixotiksテックブログ

Discussion