『ZOZOTOWN「おすすめアイテム」を支える推薦システム基盤』を支えるKubeflow実験基盤の構築と改善

開発環境 機械学習 Python Google Cloud Kubernetes Kubeflow

はじめに

こんにちは。ZOZO研究所の shikajiro です。主にZOZO研究所のバックエンド全般を担当しています。

先日のテックブログ ZOZOTOWN「おすすめアイテム」を支える推薦システム基盤 をご覧いただけたでしょうか。ZOZO研究所と連携するMLOpsチームのTJこと田島が執筆した記事なので是非御覧ください。

techblog.zozo.com

この 推薦システム基盤の推薦アルゴリズム を研究開発する際に利用した 実験基盤 の開発メンバーとして参加し、そこでAI PlatformやKubeflowを活用して効率的なML開発を試みました。今回はこの実験基盤の開発を紹介したいとおもいます。

また、推薦基盤チームのてらちゃんこと寺崎が執筆した AI Platform Pipelines (Kubeflow Pipelines)による機械学習パイプラインの構築と本番導入 はKubeflowの基本を知る上で大変参考になりますので、合わせて御覧ください。Kubeflowの説明についてはこちらの記事が充実していますので、本記事では省略しています。

techblog.zozo.com

さらに、私が半年前に執筆した 近似最近傍探索Indexを作るワークフロー を読んでいると少しだけ楽しさが増すのでぜひご覧ください。

techblog.zozo.com

目次

AI Platformの導入

プロダクション環境の推薦基盤はGCP上で動いているため、ZOZO研究所による推薦アルゴリズムの研究開発もGCP上で行いました。

ZOZO研究所の研究開発メンバーは普段AWS環境での開発に慣れており、GCPで同等のサービスを模索しました。AI Platformである程度の開発効率を見込めそうなので、AI Platform Pipelinesを使って開発を進めました。

当時はAI PlatformやKubeflowに慣れているメンバーが居なかったため、試行錯誤しながら実験基盤の開発体験の向上を行っていきました。

Kubeflowを独自構築

開発当初のAI Platform PipelinesのKubeflowはversion 0.2と古く、何度か実験を動かすと原因不明のエラーで止まったり、実験結果が表示されない事もあり大変不安定でした。

AI Platform PipelinesではブラウザでKubeflow Pipelinesの構築ができます。しかし、バージョンの追従はGCP次第であり、自由に選択できません（執筆時点ではブラウザからKubeflow Pipelines 1.0が構築可能）。

そこで、AI Platform Pipelines運用を一旦諦め、当時最新のKubeflow Pipelines 1.0をGKEに独自に構築することで安定化させました（執筆時点ではKubeflow Pipelines 1.2が最新）。

その後、Kubeflow Pipelines 1.1がリリースされたのでインストールを試みたのですが、当時のドキュメントの完成度が高くなかった こともあり、うまくいかず一旦諦めました。

ビルドフローの安定化

当初、開発者のローカルマシンでビルドしたDockerコンテナをGCRにアップロードしてパイプラインを実行していました。

latestタグを使ってしまうと予期せぬコンテナが使われてしまったり、かと言って都度異なるtagを指定するとパイプラインの中の実装を変える必要があったりと、なかなか手間がかかっていました。さらに、開発者のローカルマシンのCPUがDockerビルドにより消費してしまい、開発がし辛い状況でした。

そこで、DockerのビルドはすべてCloud Buildに変更しました。これによりローカルマシンのCPUを使うことはなくなり、一貫したビルドフローを使うことができるので、安定した開発を行うことができるようになりました。

この流れはこちらの記事を参考にしています。 cloud.google.com

パラメータを外部ファイル化

開発時は開発者の任意のタイミングで何度も何度も学習と予測を行います。Kubeflow Pipelinesは一度作成したパイプラインをWeb画面からパラメータを変えて何度も実行できることが長所の1つです。

しかし、多くのパラメータがあるパイプラインの場合、実験の度にパラメータをすべて入力するのは少々手間です。開発者からは「気軽にコマンドラインから実験したい」と要望があったので、パラメータをyamlで管理してパイプラインにyamlを渡すことで、開発者のマシンから気軽に実行できるようになりました。

ファイルで管理することにより、設定したパラメータにコメントを残したり、Git管理できたり、入力ミスを減らすことができました。

settings.yamlの例

# project 全体の設定値
project_id: hello-zozo
webhook: https://hooks.slack.com/services/hogehoge

# Kubeflow Pipelines の設定値
experiment_name: Default
run_name_base: example
gcr_image: gcr.io/{project_id}/example:{tag}

# 学習や検証・予測などで使うパラメータ
number: 100

TyperでのCLIによる簡単な実行

パイプラインを実行するPythonコードに Typer を導入しました。TyperはPythonのCLIアプリケーションを簡単に作れるライブラリです。代表的なものにargparseやClickがありますが、Typerはさらに使いやすくされたものです。

普段の実験は以下のコマンドで実行しています。

# python <パイプラインを実行するTyper実装> run <パイプラインを定義したPythonファイルがあるディレクトリ>
python pipeline.py run helloworld

複数のyamlファイルを指定し、連続して実験を行うこともできます。

python pipeline.py run helloworld --settings hoge.yaml --settings fuga.yaml

前回ビルドしたDockerをそのまま利用する場合、tag名を指定してCloud Buildをスキップすることもできます。

python pipeline.py run helloworld --tag 'docker-tag-name'

検証高速化のためのスキップ処理

パイプラインには学習パートと検証パートがあります。学習部分は前回実行して生成されたモデルを使って、検証部分だけ実装を変更して実行したい要望がありました。

Kubeflow Pipelinesには任意の位置から実行する機能はありません。新しくパイプラインを作ってはじめから実行する必要があります。

そこで、パイプラインの中にスキップフラグによる条件分岐する仕組みを追加し、生成済みのモデルを使って検証パートだけを行える仕組みを作りました。

パイプラインの実装はConditionで分岐させました。スキップする時、しない時両方のパイプラインの流れを定義することで、スキップを実現させています。

@dsl.pipeline(name="hello", description="hello world pipeline")
def pipeline(skip_build: bool, ...):
    first = dsl.ContainerOp(...)
    last = dsl.ContainerOp(...)
    with dsl.Condition(skip_build == False, name="build"):
        # とても重たい処理
        build = dsl.ContainerOp(...)
        build.after(first)
        last.after(build)
    with dsl.Condition(skip_build == True, name="skip-build"):
        last.after(first)

パイプライン管理画面のGraphを見ると、スキップされていることが分かります。

失敗談と感想

みんな大好き失敗談を紹介します。

パイプラインの無理な流用でエラーが頻出する

パイプラインの実行処理高速化、管理の利便性向上のため、次の改善を行いました。

  • 既にパイプラインがある場合はそれを使い、パイプラインコードやパラメータに追加削除が合った場合はパイプラインの新しいバージョンを作って実行する

この対応が原因でパイプライン実行時にエラーが頻出し、MLエンジニアが研究開発し辛い状況を作ってしまいました。

理由の説明の前に、Kubeflow Pipelinesで使われる3つの要素について説明します。

項目 必須 説明
Run 必須 最小実行単位。実行したワークフローをRunと表現する。
Experiment 必須 Runを束ねる存在。Pipelineを指定するか、同等のyamlを直接指定して動かす。Experimentの粒度はチーム内で取り決めるのが良さそう（アルゴリズム単位、日付単位、パラメータ単位など）。
Pipeline 任意 DAGの事。パラメータを指定すればすぐ動かせる。バージョン管理可能。繰り返し実行などする場合はPipelineが必要。

Run、Experiment、Pipelineの関係がちょっとややこしいため、注意が必要です。

  • Run

    実行の最小単位です。Pipelineの1つの実行はRunで表されます。

  • Experiment

    Runを分かりやすく分類するために名前をつけるものです。複数のRunを束ねることができます。Runには必ずExperimentが必要です。指定しない場合はDefaultになります。1つのExperimentに異なるPipelineのRunをまとめても構いません。ディレクトリに近い感覚です。

  • Pipeline

    ここが曲者です。PipelineはDAGを定義したものになります。DAGを定義したソースコードを登録すれば、Pipelineとして一覧に表示されます。Pipelineを選びExperimentを指定して実行すれば、Runが新たに生成され実行されます。しかし、ワークフローを定義したPythonコードがあれば、Pipelineを登録していなくても実行できます。

「Pipelineが無くてもPipelineが実行できます」

何を言ってるか分からないと思いますが、Kubeflow PipelinesでのRunの実行にPipelineを登録しておく必要は無いのです。ソースコードがあれば直接実行できます。

Pipelines SDKを使って実行する場合、パイプラインとソースコードを同時に指定して実行できてしまいます。その時エラーにはならず、パイプラインが謎の挙動をするため気づくまで解決が困難になります。

これらが原因でしばらくの間、パイプラインの実行がややこしく、エラーが起きがちになっていました。現在はPipelineを登録せず、ExperimentとRunだけを使って実験を行うようにし、ある程度安定したらPipelineとして登録するようにしています。

Composerと比べてどうだった？

ワークフローエンジンは他にAirflow(Composer)やDigdagがあります。それぞれ触ってみた私なりの感想を書いてみたいと思います。

その前に、MLにおいてワークフローはどうあるべきかを定義したManifest for ML in production を紹介します。

Reproducible 9ヶ月前に学習したモデルが全く同じ環境で、同じデータで再学習でき、ほぼ同じ(数%以内の差)の精度を得られるべきである

Accountable 本番で稼働しているどのモデルも、作成時のパラメータと学習データ、更に生データまでトレースできるべきである

Collaborative 他の同僚の作ったモデルを本人に聞くことなく改善でき、非同期で改善とコードやデータのマージができるべきである

Continuous 手動での作業0でモデルはデプロイできるべき。統計的にモニタリングできるべき

https://docs.google.com/presentation/d/17RWqPH8nIpwG-jID_UeZBCaQKoz4LVk1MLULrZdyNCs/edit#slide=id.g6ad50e93e5_0_59docs.google.com docs.google.com

ただワークフローを動かすのではなく、関連したデータを正しく管理できるかがMLのワークフローエンジンに求められます。

  • Airflow

    GCPではComposerという名前でマネージドサービスになっており、世界的に人気があるワークフローエンジンです。画像検索の裏側 でもComposerを使っています。少し前まではComposerで動いているAirflowのバージョンがかなり古く不安定でしたが、今は割と落ち着いています。Airflowの仕組み上、「ワークフロー実行中にDAGファイルを更新すると途中からは更新した内容で動き出してしまう」ようになっています。これはインタラクティブに開発できて自由度が高いといえば聞こえは良いですが、過去に動いたワークフローがどのバージョンのDAGで動いたか全く分からず保証も無いため、運用管理コストがとても高くなります。

  • Digdag

    GCPではマネージドサービスがなく、自前で構築する必要があります。Airflowと違ってDAGがきちんと管理されるので実行したワークフローがどのDAGで動いたのか、パラメータはなんだったのかが明白でとても扱いやすいです。しかし、最初に書きましたが自前で構築する必要があるため、初期構築・運用管理が大変になります。

  • Kubeflow

    AI PlatformとしてマネージドサービスがありDigdagと同じようにDAGがパイプラインとして管理されているので、実行したワークフローに使ったDAGやパラメータが一目瞭然です。実行がすべてk8sのpodとして動くのでCPU/GPU・メモリなどのリソース管理、スケールしやすいのもとても良いです。まだ1.0になったばかりで新機能が続々と追加されており、技術を追っていくのが大変ですが、MLの開発効率には大きく寄与しそうです。

Manifest for ML in production をふまえると、Airflowでは実現できないことが分かります。Digdagでもできそうですが、MLに特化したKubeflowはパラメータ管理に秀でておりManifestを実現するための第一の選択肢になりそうです。ただ、そもそも目的とするものが違うため「Kubeflowは良い、Airflowはだめだ」ということではないです。それぞれにメリット・デメリットがあるので見極める必要があります。

まとめ

このような実験基盤の改善を経て、研究開発を日々行い、 推薦システム基盤の推薦アルゴリズム はできあがっていっています。

「Kubeflow Pipelinesを導入すれば全部うまくいく！」ということはなく、チームメンバーで使いやすいように日々改善を行い、安定させていく事が重要になります。まだまだKubeflow Pipelinesの力を最大限発揮できていませんが、今後はより安定したワークフローエンジンとして動かせるよう、改善していきたいと思います。

さいごに

ZOZOテクノロジーズではZOZO研究所のMLエンジニア、バックエンドエンジニアのメンバーを募集しております。 https://hrmos.co/pages/zozo/jobs/0000029hrmos.co hrmos.co

カテゴリー