CUBE SUGAR CONTAINER

技術系のこと書きます。

Python: MLflow Projects を使ってみる

MLflow は MLOps に関連した OSS のひとつ。 いくつかのコンポーネントに分かれていて、それぞれを必要に応じて独立して使うことができる。 今回は、その中でも MLflow Projects というコンポーネントを使ってみる。 MLflow Projects を使うと、なるべく環境に依存しない形で、ソフトウェアのコードをどのように動かすかを示すことができる。

機械学習に限らず、特定のリポジトリのソースコードをどのように動かすかはプロジェクトによって異なる。 よくあるのは Makefile やシェルスクリプトを用意しておくパターンだと思う。 ただ、そういったスクリプトをプロジェクトごとに、環境依存のできるだけ少ない形で書いていくのはなかなかつらい作業になる。 MLflow Projects を使うと、そういった作業の負担を軽減できる可能性がある。

今回使った環境は次のとおり。

$ sw_vers
ProductName:    Mac OS X
ProductVersion: 10.15.7
BuildVersion:   19H2
$ python -V
Python 3.8.6
$ pip list | grep -i mlflow
mlflow                    1.11.0
$ docker version   
Client: Docker Engine - Community
 Cloud integration  0.1.18
 Version:           19.03.13
 API version:       1.40
 Go version:        go1.13.15
 Git commit:        4484c46d9d
 Built:             Wed Sep 16 16:58:31 2020
 OS/Arch:           darwin/amd64
 Experimental:      false

Server: Docker Engine - Community
 Engine:
  Version:          19.03.13
  API version:      1.40 (minimum version 1.12)
  Go version:       go1.13.15
  Git commit:       4484c46d9d
  Built:            Wed Sep 16 17:07:04 2020
  OS/Arch:          linux/amd64
  Experimental:     false
 containerd:
  Version:          v1.3.7
  GitCommit:        8fba4e9a7d01810a393d5d25a3621dc101981175
 runc:
  Version:          1.0.0-rc10
  GitCommit:        dc9208a3303feef5b3839f4323d9beb36df0a9dd
 docker-init:
  Version:          0.18.0
  GitCommit:        fec3683

もくじ

下準備

あらかじめ、今回使うパッケージをインストールしておく。

$ pip install mlflow

また、Python のパッケージ以外にも Docker をインストールしておこう。

$ brew cask install docker

最も単純なサンプル

MLflow Projects を最低限使うのに必要なのは MLproject という YAML ファイルを用意することだけ。 ここにプロジェクトを実行するときのエントリポイントを記述していく。 基本的には、単純にシェルで実行するコマンドを書く。

以下は wget(1) を使って Iris データセットの CSV をダウンロードする MLproject を用意している。 エントリポイントは複数書けるけどデフォルトでは main という名前が使われる。

$ cat << 'EOF' > MLproject               
name: example

entry_points:
  main:
    command: "wget -P data https://raw.githubusercontent.com/mwaskom/seaborn-data/master/iris.csv"
EOF

上記のファイルが用意できたら、実際に実行してみよう。 MLproject がある場合には mlflow コマンドの run サブコマンドが使えるようになる。 たとえば、以下のようにして実行する。 .MLproject のある場所を表していて、--no-conda は Conda の仮想環境を使わずに実行することを示す。

$ mlflow run . --no-conda
2020/10/19 19:31:56 INFO mlflow.projects.utils: === Created directory /var/folders/1f/2k50hyvd2xq2p4yg668_ywsh0000gn/T/tmp84mtkoa7 for downloading remote URIs passed to arguments of type 'path' ===
2020/10/19 19:31:56 INFO mlflow.projects.backend.local: === Running command 'wget -P data https://raw.githubusercontent.com/mwaskom/seaborn-data/master/iris.csv' in run with ID '96c013ed298f48be81ceafc9e64bb726' === 

...

2020-10-19 19:31:57 (8.16 MB/s) - `data/iris.csv' へ保存完了 [3858/3858]

2020/10/19 19:31:57 INFO mlflow.projects: === Run (ID '96c013ed298f48be81ceafc9e64bb726') succeeded ===

すると、次のようにファイルがダウンロードされることがわかる。

$ ls data
iris.csv
$ head data/iris.csv 
sepal_length,sepal_width,petal_length,petal_width,species
5.1,3.5,1.4,0.2,setosa
4.9,3.0,1.4,0.2,setosa
4.7,3.2,1.3,0.2,setosa
4.6,3.1,1.5,0.2,setosa
5.0,3.6,1.4,0.2,setosa
5.4,3.9,1.7,0.4,setosa
4.6,3.4,1.4,0.3,setosa
5.0,3.4,1.5,0.2,setosa
4.4,2.9,1.4,0.2,setosa

Docker コンテナで実行する

ただ、先ほどの例はシステムに wget(1) がインストールされていることを仮定している。 しかしながら、誰もがシステムに wget(1) をインストールしているとは限らないだろう。 かといって、各プラットフォームごとに wget(1) をインストールするスクリプトや、方法について README には書きたくない。 そこで、続いては MLflow Projects を Docker コンテナ上で実行できるようにしてみよう。

MLflow Projects を Docker コンテナ上で実行する場合には、docker_env という要素を MLproject ファイルに書く。 たとえば、次のサンプルでは example/mlflow-docker というイメージでこのプロジェクトを実行しますよ、という宣言になる。

$ cat << 'EOF' > MLproject               
name: example

docker_env:
  image: example/mlflow-docker

entry_points:
  main:
    command: "wget -P data https://raw.githubusercontent.com/mwaskom/seaborn-data/master/iris.csv"
EOF

ただ、上記を実行するにはシステムのローカルないし Docker リポジトリに example/mlflow-docker というイメージが必要になる。 そこで、たとえば次のような Dockerfile を用意してみよう。 これで、wget(1) のインストールされた Ubuntu 20.04 LTS のイメージがビルドできる。

$ cat << 'EOF' > Dockerfile
FROM ubuntu:20.04

RUN apt-get update \
 && apt-get -yq dist-upgrade \
 && apt-get install -yq \
      wget \
 && apt-get clean \
 && rm -rf /var/lib/apt/lists/*
EOF

イメージのビルド手順も設定ファイルに書きたいので Docker Compose を使うことにする。

$ cat << 'EOF' > docker-compose.yaml
version: "3.8"

services:
  mlflow:
    image: example/mlflow-docker
    build:
      context: .
EOF

これで特に何を考えずイメージがビルドできる。

$ docker-compose build
Building mlflow
Step 1/2 : FROM ubuntu:20.04

...

Successfully tagged example/mlflow-docker:latest

システムのローカルに必要なイメージが登録されたことを確認しよう。

$ docker image list | grep example
example/mlflow-docker                latest                                           09c03e33daea        3 minutes ago       111MB

動作確認の邪魔になるので、先ほどダウンロードしたファイルは一旦削除しておこう。

$ rm -rf data

それでは、満を持して MLflow Projects を実行してみよう。 今度は実行環境として Docker を使うことが設定ファイルからわかるので --no-conda オプションをつける必要はない。

$ mlflow run .
2020/10/19 19:50:16 INFO mlflow.projects.docker: === Building docker image example ===

...

2020-10-19 10:50:18 (5.87 MB/s) - 'data/iris.csv' saved [3858/3858]

2020/10/19 19:50:19 INFO mlflow.projects: === Run (ID '18f4068e82434b119dc7cf9979082f9b') succeeded ===

しめしめ、これでカレントディレクトリにファイルがダウンロードされるはず…と思っていると何も見当たらない。

$ ls
Dockerfile      docker-compose.yaml
MLproject       mlruns

なぜ、こうなるか。 MLflow Projects は、Docker コンテナを実行するときに、ホストのディレクトリをコンテナにコピーしている。 あくまで、共有ではなくコピーなので、コンテナ上にファイルをダウンロードしたからといってホストには反映されない。

この問題を回避するには、ホストのボリュームをコンテナにマウントする必要がある。 たとえば、次のようにしてホストのディレクトリをコンテナの /mnt にマウントできる。 そして、実行するコマンドでは、保存先のディレクトリを /mnt 以下にする。

$ cat << 'EOF' > MLproject               
name: example

docker_env:
  image: example/mlflow-docker
  volumes: ["$(pwd):/mnt"]

entry_points:
  main:
    command: "wget -P /mnt/data https://raw.githubusercontent.com/mwaskom/seaborn-data/master/iris.csv"
EOF

もう一度実行してみよう。

$ mlflow run .
2020/10/19 19:54:32 INFO mlflow.projects.docker: === Building docker image example ===

...

2020-10-19 10:54:34 (776 KB/s) - '/mnt/data/iris.csv' saved [3858/3858]

2020/10/19 19:54:34 INFO mlflow.projects: === Run (ID '4372f3f53dbf478bbc8ac96b4f7d5811') succeeded ===

すると、次のようにホスト側にディレクトリとファイルができる。

$ ls
Dockerfile      MLproject       data            docker-compose.yaml mlruns
$ head data/iris.csv 
sepal_length,sepal_width,petal_length,petal_width,species
5.1,3.5,1.4,0.2,setosa
4.9,3.0,1.4,0.2,setosa
4.7,3.2,1.3,0.2,setosa
4.6,3.1,1.5,0.2,setosa
5.0,3.6,1.4,0.2,setosa
5.4,3.9,1.7,0.4,setosa
4.6,3.4,1.4,0.3,setosa
5.0,3.4,1.5,0.2,setosa
4.4,2.9,1.4,0.2,setosa

エントリポイントを増やしてみる

続いては、例としてエントリポイントを増やしてみることにしよう。 example という Python のパッケージを用意する。

$ mkdir -p example
$ touch example/__init__.py

そして、その中に preprocess.py というモジュールを追加する。 これは CSV ファイルから読み込んだ Pandas の DataFrame を Pickle にして保存するコードになっている。

$ cat << 'EOF' > example/preprocess.py
#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""CSV ファイルから読み込んだ DataFrame を Pickle として永続化する"""

import os
import logging
import pathlib

import pandas as pd

LOGGER = logging.getLogger(__name__)


def main():
    logging.basicConfig(level=logging.INFO)

    # データの置き場所
    data_dir = pathlib.Path('/mnt/data')

    # CSV ファイルを読み込む
    raw_data_file = data_dir / 'iris.csv'
    LOGGER.info(f'load data from {raw_data_file}')
    df = pd.read_csv(str(raw_data_file))

    LOGGER.info(f'data shape: {df.shape}')

    # Pickle として永続化する
    parsed_data_file = data_dir / 'iris.pickle'
    LOGGER.info(f'save data to {parsed_data_file}')
    df.to_pickle(parsed_data_file)


if __name__ == '__main__':
    main()
EOF

以下のようにエントリポイントを fetchpreprocess に分離してみる。

$ cat << 'EOF' > MLproject
name: example

docker_env:
  image: example/mlflow-docker
  volumes: ["$(pwd):/mnt"]

entry_points:
  fetch:
    command: "wget -P /mnt/data https://raw.githubusercontent.com/mwaskom/seaborn-data/master/iris.csv"
  preprocess:
    command: "python3 -m example.preprocess"
EOF

先ほどのモジュールを見るとわかる通り、動作には Pandas が必要になる。 そこで、Docker イメージも次のように更新しておく。

$ cat << 'EOF' > Dockerfile
FROM ubuntu:20.04

RUN apt-get update \
 && apt-get -yq dist-upgrade \
 && apt-get install -yq \
      wget \
      python3-pip \
 && apt-get clean \
 && rm -rf /var/lib/apt/lists/* \
 && pip3 install --user pandas
EOF
$ docker-compose build

新しく追加した preprocess エントリポイントを実行してみよう。 エントリポイントを指定するときは -e オプションを使う。

$ mlflow run -e preprocess .
2020/10/19 20:11:00 INFO mlflow.projects.docker: === Building docker image example ===

...

INFO:__main__:load data from /mnt/data/iris.csv
INFO:__main__:data shape: (150, 5)
INFO:__main__:save data to /mnt/data/iris.pickle
2020/10/19 20:11:02 INFO mlflow.projects: === Run (ID '6bae8e1f7d3e4bde85fa8e9922d2938d') succeeded ===

これで、次のように Pickle フォーマットのファイルができあがる。

$ ls data       
iris.csv    iris.pickle
$ python -m pickle data/iris.pickle
     sepal_length  sepal_width  petal_length  petal_width    species
0             5.1          3.5           1.4          0.2     setosa
1             4.9          3.0           1.4          0.2     setosa
2             4.7          3.2           1.3          0.2     setosa
3             4.6          3.1           1.5          0.2     setosa
4             5.0          3.6           1.4          0.2     setosa
..            ...          ...           ...          ...        ...
145           6.7          3.0           5.2          2.3  virginica
146           6.3          2.5           5.0          1.9  virginica
147           6.5          3.0           5.2          2.0  virginica
148           6.2          3.4           5.4          2.3  virginica
149           5.9          3.0           5.1          1.8  virginica

[150 rows x 5 columns]

エントリポイントをまとめる

ただ、先ほどのように複数のエントリポイントがあると何をいつどのように実行すれば良いのかが不明瞭になる。 それでは MLflow Projects を導入した意味が薄れてしまうかもしれない。 おそらく、エントリポイントは先ほどよりももっと大きなタスクで分けた方が良いと思う。 たとえば、モデルの学習と結果の可視化のような粒度であれば、分けていても混乱は招かないかもしれない。 あまり巨大すぎるエントリポイントを作るのも、後から取り回しが悪くなるだろう。

さて、前置きが長くなったけど、ここではあくまで一例としてエントリポイントを Metaflow を使ってまとめてみよう。 まずはデータのダウンロード部分も Python のコードにしてしまう。 めんどくさければ subprocess モジュールを使って wget(1) をキックするだけでも良いと思う。

$ cat << 'EOF' > example/fetchdata.py
#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""データセットを取得して CSV として保存する"""

import logging
import pathlib
import os
from urllib.parse import urlparse

import requests

LOGGER = logging.getLogger(__name__)

DATASET_URL = 'https://raw.githubusercontent.com/mwaskom/seaborn-data/master/iris.csv'


def main():
    logging.basicConfig(level=logging.INFO)

    # データの置き場所
    data_dir = pathlib.Path('/mnt/data')

    # ディレクトリがなければ作る
    os.makedirs(str(data_dir), exist_ok=True)

    # 保存先のファイルパス
    url_path = urlparse(DATASET_URL).path
    save_filepath = data_dir / pathlib.Path(url_path).name

    # データセットとの接続を開く
    LOGGER.info(f'starting download: {DATASET_URL}')
    with requests.get(DATASET_URL, stream=True) as r:
        # 2XX でないときは例外にする
        r.raise_for_status()
        # 書き込み先のファイルを開く
        with open(save_filepath, mode='wb') as w:
            # チャンク単位で書き込んでいく
            for chunk in r.iter_content(chunk_size=8192):
                w.write(chunk)
    LOGGER.info(f'download complete: {save_filepath}')


if __name__ == '__main__':
    main()
EOF

そして、次のようにデータの取得と加工をパイプライン化してみる。 ここでは Metaflow を使ったけど、別に何を使っても良いと思う。 手軽だったので選んだに過ぎない。

$ cat << 'EOF' > entrypoint.py
#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import metaflow

from example import fetchdata
from example import preprocess


class MyFlow(metaflow.FlowSpec):

    @metaflow.step
    def start(self):
        self.next(self.fetch)

    @metaflow.step
    def fetch(self):
        fetchdata.main()
        self.next(self.preprocess)

    @metaflow.step
    def preprocess(self):
        preprocess.main()
        self.next(self.end)

    @metaflow.step
    def end(self):
        pass


if __name__ == '__main__':
    MyFlow()
EOF

あとは、次のようにパイプラインを実行するように MLproject を編集する。 説明が遅くなったけど、次のように可変な部分はパラメータとして与えることができる。

$ cat << 'EOF' > MLproject               
name: example

docker_env:
  image: example/mlflow-docker
  volumes: ["$(pwd):/mnt"]
  environment: [["USERNAME", "$(whoami)"]]

entry_points:
  main:
    parameters:
      command:
        type: string
        default: run
    command: "python3 -m entrypoint {command}"
EOF

動作に必要なパッケージが増えたので、Docker イメージも編集する。

$ cat << 'EOF' > Dockerfile
FROM ubuntu:20.04

RUN apt-get update \
 && apt-get -yq dist-upgrade \
 && apt-get install -yq \
      wget \
      python3-pip \
 && apt-get clean \
 && rm -rf /var/lib/apt/lists/* \
 && pip3 install --user pandas requests metaflow
EOF
$ docker-compose build

先ほどダウンロードして加工したファイルは動作確認の邪魔になるので一旦削除しよう。

$ rm -rf data

ディレクトリ構成は以下のようになる。

$ ls     
Dockerfile      docker-compose.yaml example
MLproject       entrypoint.py       mlruns

それでは、実行してみよう。

$ mlflow run .
2020/10/19 20:24:57 INFO mlflow.projects.docker: === Building docker image example ===

...

2020-10-19 11:25:04.898 Done!
2020/10/19 20:25:05 INFO mlflow.projects: === Run (ID 'd9dd8d4daa264db99004142b60640307') succeeded ===

これで、次のようにファイルができる。

$ ls data 
iris.csv    iris.pickle

Python のコードから実行する

ちなみに、これまでの例は MLflow Projects を mlflow run コマンドから実行していた。 もちろん Python のコードからも実行できる。

たとえば Python の REPL を起動しよう。

$ python

あとは mlflow をインポートして mlflow.projects.run() を呼ぶだけ。

>>> import mlflow
>>> mlflow.projects.run(uri='.')

いじょう。

参考

複数ステップで構成された MLflow Projects のサンプルプロジェクトが以下にある。

github.com