mercariengineering
twitter
twitter

Polyaxon を機械学習の実験プラットフォームとして導入した話

Author:
NakagawaMai
 , 
Polyaxon を機械学習の実験プラットフォームとして導入した話

Merpay Advent Calendar 2020 の12日目は、メルペイで MLOps を担当している ML Platform チームの Mai Nakagawa がお送りします。

今回は、機械学習の実験を効率よく実行・管理するために Polyaxon という OSS のツールを導入した話をします。

Polyaxon 導入前の課題

Polyaxon について言及する前に、まずは Polyaxon 導入前に抱えていた課題について話します。それは機械学習の実験フェーズが属人化してしまっていたことです。下の図のように機械学習のワークフローを大きく実験フェーズ（Experiment Phase)、本番環境に機械学習モデルをデプロイするフェーズ（Productionize Phase）、本番環境で機械学習モデルを運用フェーズ（Operating Phase）の3フェーズに分けると、Production PhaseとOperating Phaseにはツールが導入されていて属人性が排除されているものの、実験フェーズは各 ML エンジニアがそれぞれのやり方で実験を行っていて、コードが各 ML エンジニアの実験環境に依存している部分がありました。

一般的に、モデルのコードは、ソフトウェアエンジニアが普通に書くコードに比べて複雑になりがちです。また ML エンジニア全員が必ずしもソフトウェア・エンジニアリングに明るくないこともあり、担当者が異動したり退職したときに実験のコードを引き継いでメンテナンスしていくのが難しくなるリスクがあります。そのため、モデル作成が属人的になりがちという課題があります。

このような背景のもとメルペイでは、属人性や固有の環境依存を排除して実験の再現性を担保しつつ、導入が簡単で、ソフトウェア・エンジニアリングに精通していない ML エンジニアにも使いやすいツールが求められていました。また、このあたりは Polyaxon の他にも Google AI Platform, AWS Sage Maker, Kubeflow, MLflow など様々なツールが活発に開発されているため、将来的に他のツールに乗り換えることが容易かどうかも重要視していました。

Polyaxon の紹介

上記のような課題を解決するため、私たちは Polyaxon に注目しました。Mercari US の機械学習チームがすでに Polyaxon を導入していて、彼らが構築や運用のノウハウを持っていることも大きな選定理由の一つでした。

Polyaxon は kubernetes 上で Docker コンテナを使って機械学習のコードを実行するフレームワークです。主な機能には、Jupyter Notebook、ハイパーパラメータサーチ、機械学習のワークフローに特化した Dashboard などがあります。ML エンジニアは普段と同じように実験コードを書き、追加で簡単な YAML ファイルを書くだけで、これらの機能を使って効率的に実験を行うことができるようになります。

Jupyter Notebook

Jupyter Notebook を立ち上げる場合は、例えば、以下のような YAML ファイルを書くだけで、kubernetes 上の Jupyter Notebook 環境が立ち上がり、そこで実験コードを書くことができます。build.image で指定した Docker image が build.build_steps で指定したコマンドを実行したものを Jupyter Notebook 環境として Polyaxon の Web UI から利用できます:

version: 1
kind: notebook
environment:
  resources:
    cpu:
      requests: 3
      limits: 4
    memory:
      requests: 4096
      limits: 4096
build:
  image: tensorflow/tensorflow:2.0.3-py3
  build_steps:
    - pip3 install jupyter

上記のように YAML に記述するだけで、実験に必要な計算リソースを持つ Jupyer Notebook サーバーが立ち上がります。また詳しくは後述しますが、メルペイではインフラとなる Kubernetes クラスターとして GKE (Google Kubernetes Engine) を利用しています。そのためメルペイで構築した Polyaxon では GKE の node pool を指定することで、CPU core 数やメモリ容量だけでなく、GPU インスタンスや Cloud TPU といった機械学習に特化したインスタンスタイプも指定することが可能です。

また build.image で tensorflow や pytorch など公式 docker image を使うことで、環境構築の手間を省けるのもメリットです。

ハイパーパラメータチューニング

上記では Jupyter Notebook を立ち上げて、そこで実験コードを書きましたが、Polyaxon にコードだけを submit して実行させることも可能です。単一のパラメータセットを Polyaxon 上で実験することもできるし、パラメータ空間を指定して、その中で実験を回してハイパーパラメータチューニングを行うことが可能です。メルペイが利用している Polyaxon 0.6.1 では以下のチューニング手法がサポートされています:

  • Grid search
  • Random search
  • Bayesian optimization
  • Early stopping
  • Hyperband

ハイパーパラメータチューニングを行う際は、まず実行するための Python スクリプト（この例では train.py）を作成します。次に以下のような YAML ファイルを用意します。それだけで Polyaxon 上でハイパーパラメータチューニングが行えます。ハイパーパラメータ空間は、hptuning.matrix で指定します。また、hptuning.concurrency で異なるパラメータセットの実験の並列実行数を指定できます。例えば hptuning.concurrency が 10 だと、10 個の Pod（kubernetes上のJob実行単位）が作成され、並列実行に利用されます:

---
version: 1

kind: group

build:
  image: python:3.6
  build_steps:
    - pip install --no-cache-dir -r requirements.txt

declarations:
  train_size: 0.8

hptuning:
  concurrency: 10
  random_search:
    n_experiments: 200

  matrix:
    n_estimators:
      values: [50, 100, 200, 500, 1000]
    max_depth:
      values: [3, 5, 7, 9]
    max_name_ngram_range:
      values: [1, 2, 3]
    max_description_name_range:
      values: [1, 2, 3]

run:
  cmd: python -u src/train.py \
    --train_size {{ train_size }} \
    --n_estimators {{ n_estimators }} \
    --max_depth {{ max_depth }} \
    --name_ngram_range 1 {{ max_name_ngram_range }} \
    --description_ngram_range 1 {{ max_description_name_range }}

scikit-learn など機械学習のライブラリが提供する並列処理だと、CPUコア数より大きな数の並列処理は効率的に行うことができません。これに対して Polyaxonでは hptuning.concurrency で指定した数だけ kubernetes 上に pod を作って実行するので、kubernetes が提供できる pod の数だけ並列実行が可能です。

Dashboard

Polyaxonでは機械学習の実験フェーズに特化した Web UI が提供されていて、Polyaxon 上で実行した実験の結果を見ることができます。過去の実験の結果から、どのようなパラメータで実験して、どのようなメトリックスが得られたかを一覧することができます:

また各実験の結果のメトリックスをグラフなどでビジュアライズして比較する機能もあり、効率的にパラメータ探索を行うことができます:

メルペイの機械学習チームへの導入と効果

システム・アーキテクチャ

メルペイでは、インフラに GCP (Google Cloud Platform) を利用しているため、Polyaxon も GCP 上に構築しました。システム・アーキテクチャは下記のようになっています。

polyaxon-core (API サーバー、Scheduler, 検索などのコア機能) と、実験を行うときのリソースは GKE (Google Kubernetes Engine) を利用します。polyaxon-core が利用するデータの中で永続化が必要なものに DB とファイルシステムがあり、DB は Cloud SQL を、ファイルシステムは GCE 上に NFS サーバーを構築して利用しています。実験で使ったデータセット、モデル、ログなどは GCS (Google Cloud Storage) に保存しています。

GKE の中で定常的に動くのは polyaxon-core だけで、このサーバープログラムでは重い処理を行わないため、小さいインスタンスタイプの node 上で動かしています。

実験を行う際には、大きなインスタンスタイプの node や、Cloud TPU を使う必要も出てきますが、これらは別の node pool で定義しておき、これらの node pool では autoscaling を有効にした上で、node 数は 0 にしておきます。そして実験を動かすときに autoscaling により node が立ち上がり、その node 上の pod で実験を行います。そのため、普段は polyaxon-core のための node の分しか料金はかかりません。

また実験で行う node はなるべく Preemptible VM instance を使うことで実験で使う node の料金を抑えています。Preemtible VM instance を使うと、CPU のみのインスタンスで 1/5、GPU インスタンスや Cloud TPU では 1/3 に料金を抑えることができます。これにより ML エンジニアが気軽に実験を行えるようにしています。もちろん Preemtible VM instance だと途中で instance が落ちてしまったり、最長で 24 時間までしか利用できなかったりするので、ユースケースに応じて通常の VM instance も利用します。

実験を行う際の用意する YAML ファイルで GKE の node pool を指定できるので、そこで実験を行う際に利用したい Instance type に対応した node pool を指定して実験を行うことができます。このあと出てくる例では environment.node_selector.polyaxon で experiments-preemptible という node pool を指定していて、GKE 上でこの node pool が Preemptible VM instance を利用するように設定されています。

なお環境構築の面でも属人性は排除されています。環境構築には、メルカリ/メルペイの他のインフラと同様に terraform を使っています。terraform の設定ファイルはソースコードリポジトリで管理され、terraform の実行は CI ツールから自動的に行っています。GKE の node pool の設定も terraform で管理されているため、ML エンジニアの要求を満たす新しいタイプの計算リソースが必要になった際にも、そのための新しい node pool を terraform で簡単に定義できます。

メルペイでの利用例

メルペイでは主に以下のユースケースで機械学習を使っています:

  • メルペイスマート払いという与信サービスにおける与信枠の算出
  • 不正検知システム

どちらでも Polyaxon の利用が開始されていて、主に以下のようなステップで利用されています。

  1. BigQuery(BQ) クエリでトレーニングデータを作成して Google Cloud Storage(GCS) に保存する
  2. GCS に保存されたデータセットからパラメータチューニングなどを行う

メルペイでは機械学習や分析のためのデータが BigQuery に集約されていますが、BigQuery クエリで取得したデータをそのままトレーニングデータとして利用するのではなく、一度 GCS に保存して、それをトレーニングデータとして利用します。これは実験の再現性を担保するためです。というのも BigQuery に格納されているデータは時間が経過するにつれて更新されていくため、BigQuery のデータをそのまま使うと、同じパラメータセットで同じ実験コードを実行しても、タイミングによって結果が異なってきてしまいます。また実験コードが直接 BigQuery クエリを利用しない別の理由として、ハイパーパラメータチューニングの際の並列実行の数だけ BigQuery クエリが実行して、BigQuery の Slot を専有してしまって他のメンバーの業務に支障をきたすのを避ける目的もあります。

BQ クエリでトレーニングデータを作成して GCS に保存する

擬似コードによる例ですが、以下のようにして BigQuery クエリで取得したデータを pickle にして GCS に保存します:

import tempfile
import pandas as pd
from polyaxon_client.tracking import get_data_paths, Job
from polystores.stores.gcs_store import GCSStore
from polystores.stores.manager import StoreManager

df = pd.read_gbq("...")  # データセットを生成するためのクエリ
tmp_file = os.path.join(tempfile.mkdtemp(), "data.pickle")
df.to_pickle(tmp_file)

job = Job()

# 固定パスに保存
DATA_GCS_KEY = "data-gcs"
data_path = os.path.join(get_data_paths()[DATA_GCS_KEY], job.username, job.project_name)
StoreManager(store=GCSStore()).upload_file(tmp_file, data_path)

# Job ID 付きのパスに Snapshot として保存
job.outputs_store.upload_file(tmp_file)
print("Save the snapshot data at {}".format(job.outputs_store.path))

kind: job で YAML ファイルを定義して、上記のコード（data.pyという名前で保存）を実行します:

---
version: 1
kind: job
environment:
  node_selector:
    polyaxon: experiments-preemptible  # Preemptible VM instance の node pool を指定
build:
  image: python:3.6
  build_steps:
    - pip install --no-cache-dir -r requirements.txt
run:
  cmd: python -u data.py

この例では environment.node_selector で廉価な preemptible node を使うように設定して、Job を実行しています。そして固定パスと Job ID を付与したパスにトレーニングデータを保存しています。固定パスは次のトレーニングで使われます。Job ID を付与したパスは Snapshot として保存しておき、あとからトレーニングで利用したデータを確認したいときに使います。これはデータのバージョニングができる面でも有益です。標準出力は、Polyaxon の Web UI から確認できるのですが、Polyaxon が下記のように保存先の GCS に Job ID（この例では 18）を割り振ってくれます:

GCS に保存されたデータセットからパラメータチューニングなどを行う

こちらも擬似コードによる例ですが、以下のようにして pickle で保存したデータセットを GCS からダウンロードして、トレーニングと評価を行い、結果を Polyaxon に保存しています:

...

import pandas as pd
from sklearn import metrics
from polyaxon_client.tracking import Experiment, get_data_paths, get_outputs_path
from polystores.stores.gcs_store import GCSStore
from polystores.stores.manager import StoreManager

...  # ここでハイパーパラメータサーチで使うパラメータを読み込んだりする

exp = Experiment()

# pickle で保存したデータセットを GCS からダウンロードしてメモリに展開する
data_path = args.data_path
store = GCSStore()
sm = StoreManager(store=store)
DATA_GCS_KEY = "data-gcs"
gcs_data_path = os.path.join(get_data_paths()[DATA_GCS_KEY], exp.username, exp.project_name, "data.pickle")
data_path = os.path.join(tempfile.mkdtemp(), "data.pickle")
sm.download_file(gcs_data_path, data_path)
df = pd.read_pickle(data_path)

...  # ここで Training と Evaluation

# メトリックスを Polyaxon に保存
resulting_metrics = {
    'precision_macro': metrics.precision_score(y_test, predictions, average='macro'),
    'precision_micro': metrics.precision_score(y_test, predictions, average='micro'),
    'recall_macro': metrics.recall_score(y_test, predictions, average='macro'),
    'recall_micro': metrics.recall_score(y_test, predictions, average='micro'),
    'f1_macro': metrics.f1_score(y_test, predictions, average='macro'),
    'f1_micro': metrics.f1_score(y_test, predictions, average='micro'),
}
exp.log_metrics(**resulting_metrics)

# モデルを GCS に保存
outputs_path = os.path.join(get_outputs_path())
tmp_file = os.path.join(tempfile.mkdtemp(), "model.joblib")
joblib.dump(model, tmp_file)
exp.log_output(tmp_file)

print("Save model in {}".format(outputs_path))

このコードを、前述した「ハイパーパラメータチューニング」のところで説明した要領で YAML を用意して実行します。

コードの最後のところで、各パラメータセットによる実験結果のメトリックスを Polyaxon に送信し、作成されたモデルを GCS に保存しています。これにより、満足のいくメトリックスのパラメータセットを見つけるだけでなく、そのパラメータセットで作成されたモデルを GCS から取得することが可能です。

Polyaxon 導入による効果

Polyaxon の導入により、モデル作成に使った実験コードとパラメータセットや、評価結果のメトリクスなどの情報がトレースできるようになり、メルペイ ML エンジニアの実験フェーズにおける属人性を排除することができました。また環境構築や実験コードの実行がストレスなく行えるようになりました。特にハイパーパラメータチューニングは、Scikit-learnなどのAPIの並列処理を利用して長時間回していたときに比べて、はるかに早く結果が分かるようになりました（20時間→2時間といったケースも）。またローカル環境で実験を動かしていて気づいたら処理が落ちていたので実験を回し直した、といったこともなくなり、効率的に実験フェーズでモデル開発が行えるようになりました。

今後の展望

今後は Polyaxon で構築したモデルを、後続の本番環境にモデルをデプロイするパイプラインと連携するなどして、メルペイの機械学習ワークフローをよりなめらかなものにしていきたいと思っています。

最後に

メルペイ機械学習チームへの Polyaxon の導入に多大な貢献をしてくれた仲間たちに改めて感謝します。ありがとうございました！

  • Mercari US で Polyaxon を導入・運用していて、メルペイへの Polyaxon 導入を超強力に支援してくれた Yu さん
  • Polyaxon 構築時の GKE や terraform などのインフラ面をサポートしてくれたメルペイ SRE の komattaka さん
  • Polyaxon 導入に協力的で、真っ先に試してくれたメルペイ ML エンジニアの moriaki さんと matchan さん

明日のMerpay Advent Calendar 2020 執筆担当は、 ML エンジニアの yuhi さんです。引き続きお楽しみください。
  • Twitter
  • Facebook
  • linkedin
  • このエントリーをはてなブックマークに追加

Related article