Jupyter+MLflowが動作するKaggle環境をDockerでリモートサーバ上に構築する

3 min read読了の目安(約3300字

JupyterとMLflowが動作する環境をリモートサーバ上に構築する方法の備忘録です。

MLflow自体の使い方などは以下の記事が参考になるかと思います。

https://qrunch.net/@currypurin/entries/sECZEcH9OVg4K2iq

tl;dr

Dockerとdocker-composeがインストール済みの環境で以下のリポジトリのREADMEの通りに実行してください。

https://github.com/shinsuke27/kaggle-docker-with-mlflow

Kaggleの公式Dockerイメージをベースに、自分専用のKaggle環境が立ち上がり

でJupyterとMLflowにアクセスできるはずです。

構成

Dockerfile

DockerfileはJupyterコンテナ用とMLflowコンテナ用で別々に用意しました。

Dockerfile_jupyter

  • ベースイメージはKaggleの公式Dockerイメージ
  • KaggleのNotebookに合わせ、/kaggle/workingディレクトリを作成し、そこを作業ディレクトリとする
FROM gcr.io/kaggle-images/python:latest

RUN mkdir -p /kaggle/working && \
    conda install -y -c conda-forge mlflow

Dockerfile_mlflow

  • ベースイメージはPython:3.8
FROM python:3.8

RUN pip install --upgrade pip && \
    pip install mlflow

docker-compose.yml

version: "3"
services:
  jupyter:
    build:
      context: .
      dockerfile: Dockerfile_jupyter
    volumes:
      - ./:/kaggle
    working_dir: ${WORKING_DIR}
    ports:
      - "8080:8080"
    environment:
      HOME: ${WORKING_DIR}
    command: jupyter notebook --port 8080 --ip=0.0.0.0 --notebook-dir="${WORKING_DIR}" --no-browser --allow-root --NotebookApp.token=''
    restart: always

  mlflow:
    build:
      context: .
      dockerfile: Dockerfile_mlflow
    volumes:
      - ./working/mlruns:${WORKING_DIR}/mlruns
    working_dir: ${WORKING_DIR}
    depends_on:
      - jupyter
    ports:
      - "5000:5000"
    environment:
      HOME: ${WORKING_DIR}
    command: mlflow server --backend-store-uri ${WORKING_DIR}/mlruns --default-artifact-root ${WORKING_DIR}/mlruns --host 0.0.0.0 --port 5000
    restart: always

docker-compose.ymlではjupyterコンテナとmlflowコンテナを定義します。Jupyter側で作成されるmlrunsディレクトリをMLflowコンテナにマウントすることにより、両コンテナでデータを共有させています。
${WORKING_DIR}は環境変数として以下のように.envファイルに定義しています。

WORKING_DIR=/kaggle/working

ビルド

上記のファイルを準備した上で、docker-composeによりイメージのビルド、コンテナの起動を行います。

$ docker-compose up -d --build

Kaggleの公式イメージがPULL済みでない場合は結構時間がかかります(一時間くらいかかるかも)。

コンテナが立ち上がったら、サーバから一度ログアウトし、SSHし直します。

このとき、以下のようにSSHポートフォワーディングで複数のポートを指定します。

$ ssh your-remote-server -- -L 8080:localhost:8080 -L 5000:localhost:5000

ここまでできたらブラウザからhttp://localhost:8080/http://localhost:5000/を開きます。

JupyterとMLflowの画面が表示されるはずです。

  • Jupyter

  • MLflow

最後に、jupyterコンテナからmlflowコンテナにデータが送れるかを確認します。Jupyterで適当なNotebookを作成し、以下のコマンドを実行します。

set_tracking_urihttp://mlflow:5000と指定することで、送信先をlocalhostではなくmlflowコンテナに変更できます。

import os
import mlflow
mlflow.set_tracking_uri("http://mlflow:5000")

if __name__ == "__main__":
    # Log a parameter (key-value pair)
    mlflow.log_param("param1", 5)

    # Log a metric; metrics can be updated throughout the run
    mlflow.log_metric("foo", 1)
    mlflow.log_metric("foo", 2)
    mlflow.log_metric("foo", 3)

    # Log an artifact (output file)
    with open("output.txt", "w") as f:
        f.write("Hello world!")
    mlflow.log_artifact("output.txt")

MLflowに戻ると無事にデータが送られていることが確認できました。