MLflowをさくっと導入できるdocker-composeを作った
tl;dr
docker-composeを叩くだけでさくっと認証付きのMLflowサーバーを立てられるようにしました
こちらからどうぞ: ymym3412/mlflow-docker-compose
みなさん機械学習の実験をしていますか?
学習に使ったハイパーパラメーターやデータ、Train/Valデータのロス、、Testデータでの各種評価指標、これらを人手で管理しておくのは非常に大変です。
モデルの開発や比較実験に集中していると「あれ、この最高精度のモデルはどんな条件で実験したものだっけ...」となることもあり、再現性が失われてしまうことにもつながります。
この機械学習にまつわる課題を解決するひとつの枠組みが実験管理と呼ばれるもので、学習時に使用したハイパーパラメーターやTrain Loss、Test データでの評価結果などを記録して管理しておくものです。 代表的なものでいうとMLflowやcomet.ml、Catalystなどのツールがあります。
今回はその中のひとつのMLflowをdocker-composeを使って認証付きの環境をさくっと作れるようにしました。
こちらの記事では、docker-composeで立ちあがるコンテナの解説をします。立ち上げ方についてはリポジトリのREADMEをご覧ください。
※ここで紹介するのはMLflowにある機能の中でパラメータやモデルを保管しそれを可視化する機能を提供してくれるMLflow Trackingを使用したものです。
構成
docker-composeを使って以下のような構成のコンテナを立ち上げるように設定しています。
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に、モデルのような大きなデータはGCPのGoogle 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_USERNAME
とMLFLOW_TRACKING_PASSWORD
という環境変数にそれぞれBasic認証のユーザーとパスワードを設定しておくと、その値をTrackingサーバーへのリクエストのヘッダーに設定して認証を通過するようにしてくれます。
client側でこちらの値の設定を忘れないように注意しましょう。
まとめ
docker-composeを使ってMLflowの環境をさくっと構築できるスクリプトを紹介しました。
セキュアな環境を構築したいのであれば、ドメインを取得してクラウド上のVMに設定してそのVMでdocker-composeを実行してもいいですし、そこまでの環境が必要ない場合でもこちらのスクリプトを使って環境構築することが可能です。
MLflowは実験管理を行う上で必要な機能をシンプルに提供してくれているので、個人での実験管理にもとても便利だと思います。
質問や改善点などあればリポジトリへのissueやTwitterへのDMなどへお願いします。