10X のコスパ重視 MLOps - 10X Product Blog

どうも @metalunk です．
コスパ，大事ですよね？コストをある値以下に抑えたとき，どれだけパフォーマンスを発揮できるか，という話です．
10X で最初の機械学習プロダクトを作るにあたり，コスパを意識して MLOps 基盤を作ったので，それの紹介をします．

Stailer における ML の重要性
レジ前推薦
作りたかったもの
アーキテクチャ
stailer-suggest-batch の移行
組織の話
未来の話

Stailer における ML の重要性

10X のプロダクトである Stailer はネットスーパー / ネットドラッグストアのプラットフォームです．そして Stailer では推薦技術がとても重要です．
その重要さは Stailer プロダクト責任者の yamotty が 10Xが”検索”と”推薦”に心を燃やすワケ - 10X Product Blog で説明してくれています．ブログにある三行まとめによると

スーパーでの買い物体験は多量の”意思決定”で構成されています
Stailerはお店の買い物体験を補完するプロダクトです
ネットスーパーの買い物体験を支えるのは”検索”と”推薦”という技術です

そういうわけで，推薦技術は Stailer にとって重要です．

さらに，重要な ML 技術は推薦だけではありません．
商品の需要予測，配送計画のためのパラメータの推論，未知の商圏の売り上げ推定，倉庫の設計の最適化など，ML 技術を役立てられそうなキーワードがゴロゴロ転がっています．

ところで CEO yamotty のブログからも分かる通り，10X ではリーダーたちが ML の重要性を理解してくれています．
ML の重要性を説明することも ML エンジニアの仕事であるとは思いますが，そもそも理解してくれていると話が早いですし，仕事も進めやすいです．その点において 10X は ML エンジニアが働きやすい環境だと思います．

レジ前推薦

Stailer 最初の推薦プロダクトはレジ前推薦です．
お客さまが，欲しいものをひと通り追加してお会計をする前に「こちらもどうですか？」という提案をするやつです．

この機能は AB テストでリリースするため，ブログリリース時点ではまだ使えないパートナー / お客さまもいらっしゃると思いますが，バックエンドのシステムはすでに本番にリリースされている状態です．

レジ前推薦のゴールはカート追加数の増加です．
多くの仮説を検証するために，最初の推薦アルゴリズムはシンプルに，高速に作りたい思いがありました．よって，パーソナライズなしで，一日一回更新で，多腕バンディット問題の ε-greedy アルゴリズムを動かすことに決めました．
詳しくは，仮説が検証された頃に誰かがブログを書いてくれると期待しています．

作りたかったもの

システム面のゴールは，システムを安定稼働させること，アルゴリズムの検証と改善を高速に実施できるようにすることです．
さらに，Stailer における ML の重要性の章で説明した通り，将来 ML プロダクトをいくつも稼働することが期待されています．

一般的には最初の ML プロダクトでは MLOps に投資せず，いくつかの ML モデルが動いて Pipeline 管理などに問題が出始めてから，MLOps に投資するケースが多いと思います．
しかし，今回の要件と将来の計画を考え，自分たちは最初から MLOps に投資し，MLOps level 2 の状態を目指すことに決めました．
また，PdM, Data Scientist を含むチームメンバーたちは，過去の経験から MLOps の重要性を理解していました．それもこの決定の背中を押す一因になりました．

MLOps パートの実装をするのは自分ひとりであるため，できるだけ小さい労力で MLOps level 2 を実現するチャレンジとなりました．
Serving の実装も含めて，当初は3ヶ月でレジ前推薦をリリースしようとしていましたが，最終的に5ヶ月間ほどかかりました．

アーキテクチャ

Training pipeline
- Vertex Pipelines で動かす
- レジ前推薦では recommendable_stocks と most_popular_at_checkout の二つの pipeline が存在し，それぞれパートナーごとに定期的に実行される
Elasticsearch
- Training pipeline の成果物が入り，検索可能な状態にする
- パートナーごとに個別の index を作る
stailer-recommender
- 推薦用の API を提供する
- ε-greedy アルゴリズムが動く
stailer-server
- Stailer 本体のサーバ
- stailer-recommender を呼び出す

この「アーキテクチャ」章の残りでは，細かい個別の検討事項について説明します．
とても長いので興味のあるところだけ拾い読みしてください．読まない人は「組織の話」まで飛んでください．

Training pipeline の選択

一般に ML pipeline と呼ばれる仕組みは次のような機能を提供します．

Training を実行し，実行パラメータ，メタデータなどを管理する
component ごとに成果物をバージョン管理して保管する
component を再利用可能にする

これらの機能は自分たちの ML プロダクトに必須であるため，ML pipeline システムを導入します．
そして，ML pipeline システムの中でも Vertex Pipelines を使うことにしました．理由は次のとおりです．

managed サービスであること
- MLOps エンジニアがひとりであるためインフラ管理をサボれる必要があります
GCP サービスであること
- Stailer はほとんど全てのものが GCP 上で動いているため，ML pipeline も GCP で動くといろんなこと（データへのアクセス，権限管理，デプロイ）が楽になります

Vertex Pipelines で使える Kubeflow Pipelines SDK v2 は beta である点で二の足を踏みましたが，Kubeflow Pipelines を自前で立てるよりはマシと判断しました．

Introducing Kubeflow Pipelines SDK v2 | Kubeflow

Python function-based component vs Own container component

Kubeflow Pipelines SDK で component を実装する方法は2種類あります．Python function-based component と Own container component （正式名称がなさそうなので勝手に名付けました）です．

Build a pipeline | Vertex AI | Google Cloud

Python function-based component の実装方法は手軽で，component にしたい Python function にデコレータで library などを指定するだけです．

一方 Own container component は自前の (Docker) image を用意し，component.yaml ファイルに component の入出力，実行コマンドを書き，それらを pipeline から指定する方法です．
言語は Python に限られませんし，事前に build してあるため Pipeline 実行時に library install は不要です．

多少手間はかかりますが Own container component の方が自由度が高く，将来やりたいことが増えたときに対応できるように Own container component で進めることにしました．

ところが，実装を進めるうちに気付いたんですが Vertex Pipelines は Python function-based component を推している気がします．
というのも，後述する Vertex ML Metadata の吐き出しは Python function-based component しか公式サポートされていないし，Own container component の方がドキュメントが少ないと感じました．
たしかに Vertex AI は ML を民主化する（いろんな人が ML を触れるようにする）ためのプロダクトであるため，気軽に使える Python function-based component を推すのは自然に思えます．
ぼくは Own container component のサポートが手厚くなることを願っております．

Serving 用データストア

Training の成果物を保存し，Serving で参照するデータストアの要件は次の通りです．

partner_id, shop_id で絞り込み，ランダムにクエリできる
高速に取得できる
一日一回更新できる

ここには Elasticsearch を選定しました．理由は次の通りです．

すでに Stailer で Elasticsearch を使っている
要件通り高速に検索できる
index の alias 切り替えで一日一回の更新ができる

ε-greedy なら単純な KVS でも実現できると思いますが，Elasticsearch の方がはやく実装できるため，ちょっとオーバースペックかもしれませんがこの選択をしました．
それに，もしかしたら将来的にベクトル検索も Elasticsearch で実装できるかもしれません．

CI (Continuous Integration)

CI には GitHub Actions を使っています．CI の仕事は次の通りです．

ユニットテスト実行
すべての pipeline.py をコンパイルして JSON を吐き，差分があったら commit & push

pipeline.json ファイルを git repository で管理することで diff がわかる状態になりますし，開発者が local で compile する必要がなく，常に最新で正しい pipeline.json を維持することができます．

CD (Continuous Delivery)

CD にも GitHub Actions を使っています．CD の仕事は次の通りです．

Docker image の build & push
Vertex Pipelines の定期実行

CI で常に pipeline.json をコンパイルしており，component が動く Docker image も push されているため，それらを API に投げつけるだけで Vertex Pipelines の実行ができます．
具体的には pipeline.json といくつかの設定を google-cloud-aiplatform library を使って Job を作成して submit します．

job1 = PipelineJob(
   display_name=f'recommendable_stocks_{partner_id}',
   template_path=str(get_recommendable_stocks_pipeline_path(args.environment)),
   parameter_values={'partner_id': partner_id},
   project=get_stailer_ml_project_id(args.environment),
   location=location,
   enable_caching=False,
)
job1.submit(service_account=get_ml_pipeline_service_account(args.environment))

Monitoring

監視対象はいくつかありますが，ここでは Vertex Pipelines の監視について説明します．
Training pipeline は一日一回実行されているため，それが正しく動いているかの監視を Cloud Monitoring で実施しています．
Pipeline の失敗監視のクエリはこれです．

resource.type="aiplatform.googleapis.com/PipelineJob"
jsonPayload.state="PIPELINE_STATE_FAILED"

さらに，GitHub Actions の障害などによって，そもそも実行されなかったときにも気付けるように，直近25時間の Pipeline 実行数の Metric を作り，その値を監視しています．

リポジトリ構成

リポジトリ構成は悩みました．まず要件は次の通りです．

Training と Serving で共通する実装を共有したい
- Elasticsearch の client とか，前処理とか，定数とか
Stailer 本体も含めたモノレポにする準備はまだ整ってない

そして二つのリポジトリを作りました．

stailer-ml-pipelines リポジトリ
- Training pipeline, component の実装
- Elasticsearch client，定数などの共通実装
stailer-recommender リポジトリ
- 推薦システムの実装（レジ前推薦以外も入る予定）
- stailer-ml-pipelines を poetry の git ref で参照する

stailer-recommender を build するときは private repository である stailer-ml-pipelines を pull できる必要があります．そのために，SRE チームが用意してくれた GitHub Apps で時限付き token を生成して使っています．

この構成だと build するために認証が必要なのが面倒です．もしかしたら ML でひとつのリポジトリにしてもよかったかもしれません．それならモノレポの工夫がそれほど必要なさそうですし．

認証

このシステムでは認証が必要な箇所がたくさんあります．

GitHub Actions で build した Docker image を GCR に push
GitHub Actions で GCR から Docker image を pull
GitHub Actions から Vertex Pipelines job を submit
GitHub Actions から Kubernetes Deployment を作成
GitHub Actions から Ealsticsearch API key を作成
Vertex Pipelines から BigQuery にクエリ
Vertex Pipelines で SecretManager から API key を取得
stailer-recommender で SecretManager から API key を取得
stailer-recommender から Elasticsearch にクエリ

たくさんありますね...

このうち5番までは OIDC で時限付き token を取得して認証しています．キーレスですし，時限付きなため安全です．gh-oidc という Terraform provider を利用しています．

それ以外の認証も Terraform で管理された状態になっており，手動で作るインフラリソースはほぼありません．

Vertex ML Metadata

Vertex ML Metadata は Training のメタデータを管理するためのサービスです．
Vertex Pipelines の Python function-based component で Vertex ML Metadata を使う方法は Using Vertex ML Metadata with Pipelines | Google Codelabs にあります．
しかし Own container component でメタデータを出力する方法はドキュメントにありません．
GitHub 上での議論を見たり，Kubeflow Pipelines SDK の実装を読んでもわからなかったので，おそらくまだサポートされていないのだと思います．Vertex Pipelines の実装は見れないのでドキュメントがないと調査がかなり難しいです．

Training のメタデータを見れると便利ですが，いまの自分たちのユースケースでは必要不可欠ではありません．そういうわけで多少 Hacky な方法でもよいから実現することにしました．

import json
 
from kfp.v2.components.executor import Executor
from kfp.v2.components.types.artifact_types import Metrics
 
 
def write_metadata(executor_input: str, metrics_name: str, metadata: dict) -> None:
   """
   Helper function to write metadata to output_artifacts.
   This function uses a private property of kfp. So deprecations may easily happen.
 
   ## How to use this function
 
   1. Add a Metrics type output in your component.yaml.
      And pass the name (in this case, `metadata`) to this function's variable `metrics_name`
   ```
   outputs:
   - {name: metadata, type: Metrics}
   ```
 
   2. Receive an executor input as an input variable.
      And pass the str to this function's variable `executor_input`
   ```
   [
     python,
     -m,
     foo.bar,
     --executor-input,
     executorInput: null,
   ]
   ```
 
   3. Pass metadata as a key-value dict to `metadata` variable.
   """
   executor_input = json.loads(executor_input)
   executor = Executor(executor_input, lambda x: x)
   metrics: Metrics = executor._output_artifacts.get(metrics_name)
 
   if metrics is None:
       raise ValueError(f'Failed to find \'{metrics_name}\' in output artifacts. '
                        f'Please confirm that \'{metrics_name}\' is in outputs in component.yaml.')
 
   metrics.metadata = metadata
   executor._write_executor_output()

kfp の private property を直接触っているため推奨されませんしすぐに壊れるでしょうが，いまはこれしか方法がないし，壊れても困らないため，この実装を使っています．

今後もしメタデータの出力が必要になったら，BigQuery に出力する component を書くと思います．

stailer-suggest-batch の移行

stailer-suggest-batch とは検索キーワードサジェストの仕組みです．詳しくは以前書いたブログ 10X の検索を 10x したいパートII - 10X Product Blog で説明しています．
このシステムは Kubernetes の CronJob で稼働しています．BigQuery からデータを取得し，成形して Elasticsearch に index を作ります．やっていることが今回のレジ前推薦の Training pipeline と似ているため，いつか移行したいと考えていました．

しばらく考えていたら実装も移行も案外すぐにできそうと思えたのでやってみました．
気合いでガリガリコードを書いて stailer-ml-pipelines 上で検索キーワードサジェストの実装をし，3時間ほどすると Vertex Pipelines で動かせました．

この通り，慣れれば Kubeflow Pipelines の component や pipeline の実装はすぐにできます．いつか新しい PoC の実装を stailer-ml-pipelines に乗せるときも短時間でやれると期待しています．

組織の話

この12月に新たに機械学習&検索部が設立されました．
専任メンバーは部長の自分ひとりだけですが，機械学習，検索それぞれに心強い兼務のメンバーがいます．
機械学習も検索もどちらも，これから改善をする準備ができた状態で，成果をどんどん生み出せるボーナスタイムです．興味がある方はぜひ Job description を読んでみてください．

ソフトウェアエンジニア（機械学習） / 株式会社10X

ソフトウェアエンジニア（検索） / 株式会社10X

ところで，機械学習プロダクトにおいて Data Scientist と MLOps エンジニアの比率がとても重要であると感じています．
ここでの Data Scientist の仕事は，機械学習における課題の発見，モデルの開発，そして分析です．そして MLOps エンジニアは Data Scientist が作った PoC（Proof of Concept）を本番で動かすことに責任を持ちます．
この通り，ML プロダクトは Data Scientist と MLOps エンジニアが互いに協力することで作りあげられます．どちらの職種も余ることなく，採用したらその分だけプロダクトに価値を提供できることが理想です．

我々のような初期の ML チームでは，シンプルなモデルを使って仮説検証することが多いし，また MLOps 基盤や新サービスなど，実装するものが多いため，Data Scientist x1 に対して MLOps エンジニア x2 くらいがちょうどいいと思います．
よって，現在の SWE(ML) の Job description は MLOps スキル高めの人がマッチするようになっています．
これは DSE (Data Science & Engineering) 部の Data Scientist である @Kazk1018 と合意してあり，Job description や採用計画も彼と協力して作成しました．

なにが言いたいかというと，MLOps エンジニアが活躍できる状況を作っているということです．

未来の話

Stailer 初の ML プロダクトをリリースしました．そして，これから多くの ML プロダクトを動かす MLOps 基盤を作りました．
これからレジ前推薦の仮説検証を受けて精度の改善をしたり，新しい推薦を作ったり，全く別の ML プロダクトを作ったり，エキサイティングな仕事が待っています．

（個人的なお知らせ）そんなエキサイティングな状況ですが，我が家にかわいい男の子が産まれました！3月までは育休をもらって家庭のことに専念します．

かわいい男の子が産まれました！
お祝いしてくださる方はマイベビーリストから送っていただけると嬉しいです！

（結婚のときは準備が足らず普段のほしい物リストを公開してしまい，数学書が届いてしまった経験から今回はちゃんと赤ちゃん用の欲しいものリストを作りました）https://t.co/4HgqnhofUO pic.twitter.com/MPRcNUWKx6
— Ryusuke Chiba (@metalunk) December 28, 2022