やむやむもやむなし

やむやむもやむなし

自然言語処理やエンジニアリングのメモ

MLflowをさくっと導入できるdocker-composeを作った

tl;dr

  • docker-composeを叩くだけでさくっと認証付きのMLflowサーバーを立てられるようにしました

  • こちらからどうぞ: ymym3412/mlflow-docker-compose


みなさん機械学習の実験をしていますか?
学習に使ったハイパーパラメーターやデータ、Train/Valデータのロス、、Testデータでの各種評価指標、これらを人手で管理しておくのは非常に大変です。
モデルの開発や比較実験に集中していると「あれ、この最高精度のモデルはどんな条件で実験したものだっけ...」となることもあり、再現性が失われてしまうことにもつながります。

この機械学習にまつわる課題を解決するひとつの枠組みが実験管理と呼ばれるもので、学習時に使用したハイパーパラメーターやTrain Loss、Test データでの評価結果などを記録して管理しておくものです。 代表的なものでいうとMLflowcomet.mlCatalystなどのツールがあります。

今回はその中のひとつのMLflowをdocker-composeを使って認証付きの環境をさくっと作れるようにしました。
こちらの記事では、docker-composeで立ちあがるコンテナの解説をします。立ち上げ方についてはリポジトリのREADMEをご覧ください。
※ここで紹介するのはMLflowにある機能の中でパラメータやモデルを保管しそれを可視化する機能を提供してくれるMLflow Trackingを使用したものです。

構成

docker-composeを使って以下のような構成のコンテナを立ち上げるように設定しています。

f:id:ymym3412:20191228202232p:plain

docker-compose.yaml

version: '3'
services:
  waitfordb:
    image: dadarek/wait-for-dependencies
    depends_on:
      - postgresql
    command: postgresql:5432

  postgresql:
    image: postgres:10.5
    container_name: postgresql
    ports:
      - 5432:5432
    environment:
      POSTGRES_USER: ${POSTGRES_USER}
      POSTGRES_PASSWORD: ${POSTGRES_PASSWORD}
      POSTGRES_DB: mlflow-db
      POSTGRES_INITDB_ARGS: "--encoding=UTF-8"
    hostname: postgresql
    restart: always

  mlflow:
    build: .
    container_name: mlflow
    expose:
      - 80
      - 443
    depends_on:
      - postgresql
      - waitfordb
    volumes:
      - ${CREDENTIALS_PATH}:/opt/application_default_credentials.json
    environment:
      DB_URI: postgresql+psycopg2://${POSTGRES_USER}:${POSTGRES_PASSWORD}@postgresql:5432/mlflow-db
      GCP_STORAGE_BUCKET: "${GCP_STORAGE_BUCKET}"
      VIRTUAL_HOST: ${HOST}
      VIRTUAL_PORT: 80
      LETSENCRYPT_HOST: ${HOST}
      LETSENCRYPT_EMAIL: example@gmail.com
      GOOGLE_APPLICATION_CREDENTIALS: /opt/application_default_credentials.json
      GCLOUD_PROJECT: ${GCLOUD_PROJECT}

  nginx-proxy:
    image: jwilder/nginx-proxy
    container_name: nginx-proxy
    restart: always
    ports:
      - "80:80"
      - "443:443"
    volumes:
      - ./${HOST}:/etc/nginx/htpasswd/${HOST}
      - html:/usr/share/nginx/html
      - dhparam:/etc/nginx/dhparam
      - vhost:/etc/nginx/vhost.d
      - certs:/etc/nginx/certs:ro
      - /var/run/docker.sock:/tmp/docker.sock:ro
      - conf:/etc/nginx/conf.d
    environment: 
      DEFAULT_HOST: ${HOST}
      DHPARAM_GENERATION: "false"
      HTTPS_METHOD: noredirect
    labels:
      - "com.github.jrcs.letsencrypt_nginx_proxy_companion.nginx_proxy"


  letsencrypt-nginx-proxy-companion:
    image: jrcs/letsencrypt-nginx-proxy-companion
    container_name: nginx-proxy-lets-encrypt
    restart: always
    depends_on:
      - nginx-proxy
    volumes:
      - conf:/etc/nginx/conf.d
      - certs:/etc/nginx/certs:rw
      - vhost:/etc/nginx/vhost.d
      - html:/usr/share/nginx/html
      - /var/run/docker.sock:/var/run/docker.sock:ro
    environment:
      NGINX_PROXY_CONTAINER: nginx-proxy


volumes:
  certs:
  html:
  vhost:
  dhparam:
  conf:

ひとつひとつ紹介していきます。

nginx-proxy + letsencrypt-nginx-proxy-companion

MLflowのserverへの通信をプロキシするNGINXのコンテナです。
MLflow自体は認証の仕組みを提供していないため、NGINXやApache httpdをリバースプロキシとして使用し、そちらで認証を行うことを推奨しています。

https://www.mlflow.org/docs/latest/tracking.html#logging-to-a-tracking-server

今回はdocker-composeで簡単にNGINXのリバースプロキシを立てられるnginx-proxyを使用しています。

MLflowはリバースプロキシでBasic認証かBearer認証を行う想定で設計されていますが、HTTPで通信を行うと認証情報が盗まれてしまう可能性があります。
なのでインターネットを経由してMLflowのサーバーにアクセスするのであればHTTPSでの通信にするのが望ましいと思います。
nginx-proxyと組み合わせて、Let's Encryptの証明書発行を自動的に行ってくれるのがletsencrypt-nginx-proxy-companionです。
これはグローバルに公開されたドメインを持っていれば、そのドメインに対する証明書の発行を自動で行ってくれます。

ですがHTTPSは必須ではないので、ドメインを持っていない場合やクローズドなネットワークといった状況でも今回のdocker-composeはお使い頂けます。

このdocker-composeではBasic認証による認証を行っています。
nginx-proxyでは動作させるホスト名と同じ名称のhtpasswdファイルを作成して、それをnginx-proxyの /etc/nginx/htpasswd/ 以下に配置することでBasi認証をかけることができます。
もしドメインを持っているならドメイン名を、持っていなければmlflow.com のように適当な値を設定すればOKです。docker-composeではその値を ${HOST} に設定して使用しています。
例えば、mlflow.com というホスト名で、demo/demo-userというユーザーでBasic認証をかけたいのであれば以下のコマンドで生成されるファイルをnginx-proxyの/etc/nginx/htpasswd/に配置すればOKです。
(実際はこのファイルをローカルに置いてdocker-compose upすれば自動で配置されます)

$ sudo echo "demo:$(openssl passwd -apr1 demo-user)" >> mlflow.com
  nginx-proxy:
    image: jwilder/nginx-proxy
    container_name: nginx-proxy
    restart: always
    ports:
      - "80:80"
      - "443:443"
    volumes:
      - ./${HOST}:/etc/nginx/htpasswd/${HOST}
      - html:/usr/share/nginx/html
      - dhparam:/etc/nginx/dhparam
      - vhost:/etc/nginx/vhost.d
      - certs:/etc/nginx/certs:ro
      - /var/run/docker.sock:/tmp/docker.sock:ro
      - conf:/etc/nginx/conf.d
    environment: 
      DEFAULT_HOST: ${HOST}
      DHPARAM_GENERATION: "false"
      HTTPS_METHOD: noredirect
    labels:
      - "com.github.jrcs.letsencrypt_nginx_proxy_companion.nginx_proxy"


  letsencrypt-nginx-proxy-companion:
    image: jrcs/letsencrypt-nginx-proxy-companion
    container_name: nginx-proxy-lets-encrypt
    restart: always
    depends_on:
      - nginx-proxy
    volumes:
      - conf:/etc/nginx/conf.d
      - certs:/etc/nginx/certs:rw
      - vhost:/etc/nginx/vhost.d
      - html:/usr/share/nginx/html
      - /var/run/docker.sock:/var/run/docker.sock:ro
    environment:
      NGINX_PROXY_CONTAINER: nginx-proxy

postgresql

MLflowは各実験でのハイパーパラメータや学習時のログ、テスト時の評価指標などを保存しており、その保存先としてローカルストレージかRelational Databaseを指定できます。
ここでは保存先としてPostgreSQLを使用しています。
PostgreSQLでは${POSTGRES_USER}/${POSTGRES_PASSWORD}というDBユーザーを作成し、 そのユーザーでmlflow-db というデータベースを作成しそこにデータを保存しています。

  postgresql:
    image: postgres:10.5
    container_name: postgresql
    ports:
      - 5432:5432
    environment:
      POSTGRES_USER: ${POSTGRES_USER}
      POSTGRES_PASSWORD: ${POSTGRES_PASSWORD}
      POSTGRES_DB: mlflow-db
      POSTGRES_INITDB_ARGS: "--encoding=UTF-8"
    hostname: postgresql
    restart: always

mlflow

MLflow Trackingのserverを動かすコンテナです。
個々の実験データは前述のPostgreSQLに、モデルのような大きなデータはGCPGoogle Cloud Storageに保存するように設定しています。
Cloud Storageへ保存するために、${CREDENTIALS_PATH}でcredentialのパスを、 ${GCLOUD_PROJECT}GCPプロジェクトを、${GCP_STORAGE_BUCKET}で保存するバケット名を指定するようにしています。

  mlflow:
    build: .
    container_name: mlflow
    expose:
      - 80
      - 443
    depends_on:
      - postgresql
      - waitfordb
    volumes:
      - ${CREDENTIALS_PATH}:/opt/application_default_credentials.json
    environment:
      DB_URI: postgresql+psycopg2://${POSTGRES_USER}:${POSTGRES_PASSWORD}@postgresql:5432/mlflow-db
      GCP_STORAGE_BUCKET: "${GCP_STORAGE_BUCKET}"
      VIRTUAL_HOST: ${HOST}
      VIRTUAL_PORT: 80
      LETSENCRYPT_HOST: ${HOST}
      LETSENCRYPT_EMAIL: example@gmail.com
      GOOGLE_APPLICATION_CREDENTIALS: /opt/application_default_credentials.json
      GCLOUD_PROJECT: ${GCLOUD_PROJECT}

以上が、docker-composeでデプロイされるコンテナの説明です。

Client

機械学習のコードの中で、ハイパーパラメータやログなどをMLflowに保存するには、MLflowのclientを使ってMLflow Trackingのサーバーにデータを送る必要があります。
MLflowでは MLFLOW_TRACKING_USERNAMEMLFLOW_TRACKING_PASSWORDという環境変数にそれぞれBasic認証のユーザーとパスワードを設定しておくと、その値をTrackingサーバーへのリクエストのヘッダーに設定して認証を通過するようにしてくれます。
client側でこちらの値の設定を忘れないように注意しましょう。

まとめ

docker-composeを使ってMLflowの環境をさくっと構築できるスクリプトを紹介しました。
セキュアな環境を構築したいのであれば、ドメインを取得してクラウド上のVMに設定してそのVMでdocker-composeを実行してもいいですし、そこまでの環境が必要ない場合でもこちらのスクリプトを使って環境構築することが可能です。
MLflowは実験管理を行う上で必要な機能をシンプルに提供してくれているので、個人での実験管理にもとても便利だと思います。

質問や改善点などあればリポジトリへのissueやTwitterへのDMなどへお願いします。

参考文献