Zenn
😖

MLflowをオンプレサーバーで構築する

2022/09/05に公開
Docker
機械学習
MLflow
MLOps
tech

こんにちは．今回は実験管理ツールであるMLflowをオンプレミスで構築する方法についてご紹介します．

導入の動機

機械学習を扱っていると，ハイパーパラメータやモデル，データセットを変更しながら様々な実験を行うことになります．この際，結果の比較を効率的に行える実験管理ツールを導入することで，モデルの開発に集中できます．

この手の実験管理ツールは，MLOpsが盛んな今，TensorboardやWeight & Bias(wandb)など，様々ありますが，最近はMLflowが選択肢として大きく盛り上がっています．

しかし，様々ある実験管理ツールの中で，データの外部送信を行わずに扱える「オンプレミス」「自前サーバ」という点に着目すると選択肢はさほど多くなく，最も使われているであろうMLflowを構築してみようと思います．

MLflowとは

Mlflowはオープンソースの機械学習のライフサイクル管理ツールで，
MLflow Tracking, Projects, Models, Registoryからなるモジュールに分かれています．

今回は主に実験管理を行うMLflow Tracking Serverをオンプレのサーバに建てる方法をまとめます．
MLflow Trackingでは機械学習モデルのハイパーパラメータや精度，ログファイルをMLflow側で一括管理することができるので，実験同士の比較を簡単に行うことができます．

構成図

今回はDocker Composeを用いて，MLflow Tracking Serverを構築していきます．

https://mlflow.org/docs/latest/tracking.html#scenario-5-mlflow-tracking-server-enabled-with-proxied-artifact-storage-access より引用

MLflowでは，実験の結果等の数値だけでなく，configファイルやソースコード，学習済みモデルファイル等のバイナリファイルもStoreすることができます．この際，バイナリファイルはartifactと呼ばれ，数値とは別の場所に保存されます．

今回はMLflowのドキュメントにもあるように，実験のパラメータ等をMySQLのDBに保存し，artifactをDockerで立てたAWS S3互換のオブジェクトファイルストレージであるMinIOに保存していきます．

また，MLflowに認証をかけるためにDockerのnginx-proxyを用いて，Basic認証を導入します．

Docker Compose

というわけで，Docker Composeで以下のコンテナを立てて構築していきます

  • Nginx Proxy(jwilder/nginx-proxy)
  • MLflow Tracking Server
  • MySQL 8.0
  • MinIO

実際に動かせるものは，こちらのGithubにあります．

https://github.com/mjun0812/MLflow-Docker

まず，MLflow Tracking ServerのDockerfileは以下です．

FROM python:3.10
RUN pip install -U pip && \
    pip install --no-cache-dir mlflow mysqlclient boto3

全体のdocker-compose.ymlは以下です．下記の例ではmlflowを5000番ポートでListenさせています．

version: "3.8"
services:
  # MLflow
  mlflow:
    container_name: mlflow
    build:
      context: .
      dockerfile: Dockerfile
    expose:
      - 80
    restart: unless-stopped
    depends_on:
      - db
      - minio
    environment:
      TZ: Asia/Tokyo
      VIRTUAL_HOST: "example.com,localhost"
      MLFLOW_S3_ENDPOINT_URL: http://minio:9000
      AWS_ACCESS_KEY_ID: minio
      AWS_SECRET_ACCESS_KEY: minio
    command: >
      mlflow server 
      --backend-store-uri 'mysql+mysqldb://mlflow:mlflow@db:3306/mlflow'
      --artifacts-destination 's3://mlflow/artifacts' 
      --serve-artifacts --host 0.0.0.0 --port 80
    networks:
      - mlflow-net

  db:
    image: mysql:8.0.29
    restart: unless-stopped
    environment:
      MYSQL_USER: mlflow
      MYSQL_PASSWORD: mlflow
      MYSQL_ROOT_PASSWORD: mlflow
      MYSQL_DATABASE: mlflow
      TZ: Asia/Tokyo
    cap_add:
      # https://github.com/docker-library/mysql/issues/422
      - SYS_NICE
    volumes:
      - ./mysql/data:/var/lib/mysql
      - ./mysql/my.cnf:/etc/mysql/conf.d/my.cnf
    networks:
      - mlflow-net

  # S3互換のストレージ
  minio:
    image: minio/minio
    restart: unless-stopped
    volumes:
      - ./minio:/export
    environment:
      MINIO_ACCESS_KEY: minio
      MINIO_SECRET_KEY: minio
    command: server /export
    networks:
      - mlflow-net

  # minioコンテナ起動時にデフォルトのバケットを自動作成する
  defaultbucket:
    image: minio/mc
    depends_on:
      - minio
    entrypoint: >
      /bin/sh -c "
      until (/usr/bin/mc config host add minio http://minio:9000 minio minio) do echo 'try to create buckets...' && sleep 1; done;
      /usr/bin/mc mb minio/mlflow;
      /usr/bin/mc policy download minio/mlflow;
      exit 0;
      "
    networks:
      - mlflow-net

  nginx-proxy:
    image: jwilder/nginx-proxy
    restart: unless-stopped
    ports:
      - "5000:80"
    volumes:
      - ./nginx/htpasswd:/etc/nginx/htpasswd
      - ./nginx/certs/:/etc/nginx/certs
      - ./nginx/conf.d/proxy.conf:/etc/nginx/conf.d/proxy.conf
      - /var/run/docker.sock:/tmp/docker.sock:ro
    networks:
      - mlflow-net
networks:
  mlflow-net:
    driver: bridge

以下，docker-compose.ymlを順に見ていきます．

MLflow

下記では，MLflow Tracking Serverを立てています．
今回は同じコンテナネットワークにあるMinIOをartifactのストレージとするため，
AWS_ACCESS_KEY_IDとAWS_SECRET_ACCESS_KEYを設定しておきます．

また，ssh port forward等でMLflowのUIにアクセスすることも考えて，VIRTUAL_HOSTには
サーバーのホスト名だけでなく，localhostも追加しておきます．ポートはnginx-proxyの都合上，
80番でListenさせておきます．

mlflowのコマンドのオプションの--backend-store-urlはSQL ALchemyと同じ方法でuriを渡します．

  mlflow:
    container_name: mlflow
    build:
      context: .
      dockerfile: Dockerfile
    expose:
      - 80
    restart: unless-stopped
    depends_on:
      - db
      - minio
    environment:
      TZ: Asia/Tokyo
      VIRTUAL_HOST: "example.com,localhost"
      MLFLOW_S3_ENDPOINT_URL: http://minio:9000
      AWS_ACCESS_KEY_ID: minio
      AWS_SECRET_ACCESS_KEY: minio
    command: >
      mlflow server 
      --backend-store-uri 'mysql+mysqldb://mlflow:mlflow@db:3306/mlflow'
      --artifacts-destination 's3://mlflow/artifacts' 
      --serve-artifacts --host 0.0.0.0 --port 80
    networks:
      - mlflow-net

  db:
    image: mysql:8.0.29
    restart: unless-stopped
    environment:
      MYSQL_USER: mlflow
      MYSQL_PASSWORD: mlflow
      MYSQL_ROOT_PASSWORD: mlflow
      MYSQL_DATABASE: mlflow
      TZ: Asia/Tokyo
    cap_add:
      # https://github.com/docker-library/mysql/issues/422
      - SYS_NICE
    volumes:
      - ./mysql/data:/var/lib/mysql
      - ./mysql/my.cnf:/etc/mysql/conf.d/my.cnf
    networks:
      - mlflow-net

そして，MySQLの設定をコンテナ内の/etc/mysql/conf.d/my.cnfに置いておきます．
タイムアウトや文字コードの設定が付与されています．

[mysql]
default-character-set=utf8mb4

[mysqld]
character-set-server=utf8mb4
collation-server=utf8mb4_unicode_ci
wait_timeout = 31536000
interactive_timeout = 31536000

MinIO

AWS S3互換のオブジェクトファイルストレージである，MinIOは以下の方法で建てることができます．

ここで設定したMINIO_ACCESS_KEYとMINIO_SECRET_KEYはMLflowのAWSのKeyと同じになるように設定しましょう．

defaultbucketでは，コンテナ起動時にMLflowの保存先バケットを作成するコマンドです．
ここのentorypoint内にも，MINIO_ACCESS_KEYとMINIO_SECRET_KEYを入力するところがあるので注意してください．

  # S3互換のストレージ
  minio:
    image: minio/minio
    restart: unless-stopped
    volumes:
      - ./minio:/export
    environment:
      MINIO_ACCESS_KEY: minio
      MINIO_SECRET_KEY: minio
    command: server /export
    networks:
      - mlflow-net

  # minioコンテナ起動時にデフォルトのバケットを自動作成する
  defaultbucket:
    image: minio/mc
    depends_on:
      - minio
    entrypoint: >
      /bin/sh -c "
      until (/usr/bin/mc config host add minio http://minio:9000 minio minio) do echo 'try to create buckets...' && sleep 1; done;
      /usr/bin/mc mb minio/mlflow;
      /usr/bin/mc policy download minio/mlflow;
      exit 0;
      "
    networks:
      - mlflow-net

nginx-proxy

ここでは，MLflowにBasic認証をかけるために，nginx-proxyを立てておきます．
以下の例ではサーバーのポート5000でMLflowに接続することができます．

  nginx-proxy:
    image: jwilder/nginx-proxy
    restart: unless-stopped
    ports:
      - "5000:80"
    volumes:
      - ./nginx/htpasswd:/etc/nginx/htpasswd
      - ./nginx/certs/:/etc/nginx/certs
      - ./nginx/conf.d/proxy.conf:/etc/nginx/conf.d/proxy.conf
      - /var/run/docker.sock:/tmp/docker.sock:ro
    networks:
      - mlflow-net

今回はNginx ProxyでBasic認証をかけるため，Basic認証のuserとpasswordも用意しておきます．

# ホストPC側
cd nginx/htpasswd
htpasswd -c [domain名] [username]

jwilder/nginx-proxyでは，コンテナ内の/etc/nginx/htpasswdにドメイン名でファイルを置くことでBasic認証を付与することができます．

また，デフォルトではnginx-proxyを介したファイルの転送には容量制限があるので，これを設定で増量しておきます．
設定をコンテナ内の/etc/nginx/conf.d/proxy.confに置きます．

# /etc/nginx/conf.d/proxy.conf
client_max_body_size 5g;

以上でサーバーの構築は完了です．次に，実験を行うPythonスクリプト側の呼び出しを見てみましょう．

実験スクリプト

MLflowはver. 1.24.0から，MLflowをS3にアクセスするためのProxyとすることで，artifactの保存先のS3ストレージのcredentialをクライアント側で保持する必要がなくなりました．利便的にもセキュリティ的にも楽です．
そのため，クライアント側で持つべき情報は

  • MLflow ServerのURL
  • Basic認証のusername, password

のみとなります．URLはスクリプト内で設定し，Basic認証の情報は環境変数に設定することでスクリプトからもServerにアクセスすることができます．

import mlflow
import os

MLFLOW_TRACKING_URI="http://example.com:5000"
MLFLOW_TRACKING_USERNAME="mlflow"
MLFLOW_TRACKING_PASSWORD="mlflow"

os.environ("MLFLOW_TRACKING_USERNAME") = MLFLOW_TRACKING_USERNAME
os.environ("MLFLOW_TRACKING_PASSWORD") = MLFLOW_TRACKING_PASSWORD

mlflow.set_tracking_uri(MLFLOW_TRACKING_URI)

あとは通常通りMLflowを実行するだけです．

以上，MLflowを完全オンプレで実行する方法でした．

参考

https://qiita.com/MasafumiTsuyuki/items/9e03e285d4b9e0c41a7c

https://blog.amedama.jp/entry/mlflow-artifact-proxy

Discussion