Vertexで3ヶ月で作る運用可能なML API基盤

こんにちは。CADDiのAI LabでMLOpsエンジニアをやっている中村遵介です。

MLOpsチームは今から3ヶ月前に立ち上がったばかりの新しいチームなのですが、その前身としてAPI基盤を作っていた時期があったので、そこで得られた知見を書いていこうと思います。

背景

CADDiのAI Labは2021年の12月に立ち上がった今月1才になったばかりの組織です。その若さにも関わらず、日々有用なMLモデルが作成されていっています。

そのような中で、「新しく作ったMLモデルを素早くユーザにデリバリーしたい」という話が上がるようになりました。ここでいうユーザとはCADDi社員や社内システム、公開アプリケーションなどを指します。

そのため、AI Lab内で簡単に使用できるAPI基盤を作成することにしました。具体的には以下の体験を作ることを目指しました。

開発者に提供するAPIデプロイ体験

  • 推論コード部分だけを記述すればAPIサーバのためのDockerイメージを完成できる
  • APIに必要なテストやドキュメントは一定のテンプレートに沿って記述できる
  • CIの中でコードから簡単にAPIをデプロイできる
  • MLモデルのデプロイに必要なテスト類は簡単なCLIを通して実行することができる

ユーザに提供するAPI使用体験

  • APIの使い方やスキーマをドキュメント上で確認できる
  • 特定のエンドポイントに対してHTTPSリクエストを送ることで期待のレスポンスを得ることができる
  • API呼び出しには認可を必要とする

開発したもの

開発したものは大きく次の5つです

  • API基盤の中心であるAPIサーバ(エンドポイント)
  • APIサーバにAPIをデプロイするためのデプロイツール
  • APIの仕様を公開するためのドキュメントツール
  • APIの品質を確認するためのテストツール
  • APIの状態を監視するためのモニタリングツール

それぞれについて選定した技術に関してご紹介していこうと思います。

API基盤の概念図

APIサーバ

APIサーバ自体はGoogle CloudのVertex AI Endpointsをそのまま使用しました。Vertex AI Endpointsを簡単に説明すると「特定の条件を満たしたDockerコンテナをAPIサーバとしてホスティングするフルマネージドサービス」です。

一般的な使い方としては、

  1. 特定の条件を満たしたDockerコンテナをArtifact Registry か Container Registry にアップロードする
  2. アップロードしたDockerコンテナを、Vertex AI Model Registryにモデルとしてインポートする
  3. インポートしたモデルを、Vertex AI Endpointsにエンドポイントとしてデプロイする

と言う流れになります。

我々のチームでVertex AI Endpointsを採択した理由は、大きく2つあります。

1つ目の理由は、GPUを簡単に扱えることです。現在利用しているMLモデルは特に軽量化をしておらず、高速な推論のためにGPUを必要とするものがあります。しかし、CloudRunではGPUを扱うことができず、GKEでは比較的大きめの運用コストを見込む必要があります。 一方で、Vertex AI EndpointsはGPUをアタッチするかどうかを選択するだけで簡単にGPU環境にデプロイ可能です。ただし、デプロイするDockerイメージの中でGPUを使用するコードを書く必要はあり、アタッチできるGPUの種類はデプロイするregionによって異なるため注意が必要です。

2つ目の理由は、Vertex AI Endpointsの要求が少なかったことです。Vertex AI Endpointsは、Vertex AI Model Registryというコンテナ置き場に置いたDockerイメージをデプロイすることでサービングを開始できます。そして、Dockerイメージは以下の要件さえ満たせば自由に実装することができます。(参考 コンテナイメージの要件

  • HTTPサーバが実行されていること
  • ヘルスチェック用のエンドポイントが用意されていること
  • 予測用のエンドポイントが用意されていること
  • リクエストサイズは最大1.5MB以下であること
  • 入出力のスキーマに一定の制約を設けること

Vertex AIのカスタムトレーニングを使用してDockerイメージを作成する場合は、自動的に上の要件を満たしたDockerイメージを作成することができます。 また、独自の実験環境でモデルをトレーニングした場合でも、上記の要件を満たすようなDockerイメージを作成すれば、APIサーバとして公開できます。この手軽さからVertex AI EndpointsをAPI基盤に選択しました。

一方で、要求が少ないということは実装の自由度が高いことを意味します。Dockerイメージを作るためには、自分達自身で上記の条件を満たしたAPIサーバを正しく実装する必要があります。 要求される知識や実装コストは決して低くありません。MLモデルをデプロイする際に、ML以外の関心ごとに多くの時間を取られるのは避けたいところです。

そこで、Dockerイメージの作成には敢えて制約をかけ、TorchServeを使用することに決めました。TorchServeはPyTorchでトレーニングしたモデルをHTTPサーバから簡単に提供できるようにするツールです。 モデルの実際の推論処理をHandlerと呼ばれるクラスの中に実装する必要がありますが、それ以外はほとんど手を加える必要がありません。

Vertex Endpointsを使用した簡易なAPIインフラ図

TorchServeでAPIサーバを起動させるには、前準備として、Torch model archiverというツールを使ってHandlerの定義ファイル(handler.py)とパラメータ(model.pt)を1つのmarと呼ばれるファイルに固める必要があります。

torch-model-archiver \
    --model-name ${model-name} \
    --version 1.0 \
    --serialized-file model.pt \
    --handler handler.py \
    --extra-files src \
    --export-path model-store

固めたmarファイルをTorchServeに渡してあげることでAPIサーバが立ち上がります。APIサーバの細かな挙動は設定ファイル(ts_config.properties)を記述することで制御できます。

torchserve \
    --start \
    --ts-config=ts_config.properties \
    --models ${model-name}.mar \
    --model-store model-store

先ほどのVertex AI Model Registryの要件と非常に相性が良く、ヘルスチェックと予測用のエンドポイントが自動で作成されます。入出力のスキーマに関しても一致しています。さらにロギングや動的なバッチ処理など多くの重要な機能を自動で提供してくれます。Vertex AI EndpointsでもHTTPサーバの作成方法の1つとしてTorchServeが挙げられていました(参考リンク: Google Cloud 上の PyTorch: Vertex AI に PyTorch モデルをデプロイする方法)。 TorchServeは、予め環境が構築されたDocker imageがCPU環境/GPU環境の両方で公開されています。

ただし、TorchServeとVertex AI Endpointsを組み合わせるためには以下に気をつける必要があります。

  • CPU環境/GPU環境の両方で動かすために、それぞれの環境用にptファイル(モデルパラメータ)が必要である
  • APIに設定ファイルや他のPythonスクリプト等の追加ファイルが必要な場合、mar作成時に渡す必要があるが、1つのディレクトリにフラットに再配置されるためファイルパスを2種類記述して切り替えるなどで対応する必要がある

特に2つ目の方は注意が必要です。 torch-model-archiverは渡されたファイルを1つのmarファイルに固め、TorchServeはそれを受け取ってテンポラリディレクトリの中に展開します。この際、元のディレクトリ構成は無視されて全てのファイルがテンポラリディレクトリの直下に展開されます(参考イシュー)。 複雑な構成の場合は予めディレクトリごとzipに固めるなどの回避策を取るそうですが、我々は数ファイル程度の依存しかなかったため、スクリプトの中でインポート先を動的に切り替えて対応しました。

デプロイツール

CI/CD

デプロイに関しては初期開発段階は手動で行っていました。開発 、検証、本番の3環境を用意して、それぞれについて権限を持っている人がGoogle Cloudコンソール上もしくはgcloud CLIを通してデプロイを行なっていました。 しかし、多くの先例が示す通り手動デプロイはミスやデプロイサイクルの速度低下に繋がったため、一連の操作を1つのスクリプトにまとめ、CIからデプロイするように変更しました。

デプロイ手段として、スクリプト以外にもterraformがよく選択されます。しかし、Vertex AI Endpoints自体はterraform管理下に置くこともできても、参照先となるVertex AI Model Registryは2022年12月1日時点だとterraform記述はできませんでした。 そのため、Vertex AI Endpoints自体もterraform管理下とせず、1つのスクリプトを実行することでDockerイメージビルドからデプロイまでの全ての流れが完了するようにしました。

Vertex Endpointsデプロイまでのフローチャート

注意点として、初回のデプロイ時と2回目移行のデプロイではコマンドの中身が異なる箇所があります。Vertex AI Model RegistryとVertex AI Endpointsについては、初回実行時はリソースを作る必要があり、2回目移行は作ったリソースに紐づける形でコンテナやモデルをデプロイしていく必要があります。

# モデル名でVertex AI Model Registry内を検索。存在すればモデルIDが入り、存在しなければ空文字となる
model_id=$(gcloud ai models list \
          --project ${project_id} \
          --region ${region} \
          --filter=displayName=${model_name} \
          --format="value(name)")

# 新規にモデルを作成
if [ -z "${model_id}" ]; then
          gcloud ai models upload …….
# 既存のモデルに紐づけて新バージョンとして作成
else
          parent_model=$(gcloud ai models describe ${model_id} \
            --project ${project_id} \
            --region ${region} \
            --format="value(name)")
          gcloud ai models upload ……. –parent-model=${parent_model}
fi

また、別の注意点としてトラフィック分割の問題が挙げられます。Vertex AI Endpointsでは新しくモデルをデプロイすると、既存モデルへのトラフィックが自動的に0%になり、新規モデルへのトラフィックが100%になります。開発環境や検証環境はデプロイと同時に新規モデルに切り替わっても問題ないのですが、本番環境に関しては安全に倒して、新規モデルへのトラフィックは0%にした状態でデプロイされるようにしました。

# 指定のエンドポイントのトラフィックを取得し、mode_1=割合,model_2=割合のように並べる
current_traffic=$(gcloud ai endpoints describe ${endpoint_id} \
          --project=${project_id} \
          --region=${region} \
          --format=json | jq '.trafficSplit’)
formatted_current_traffic=$(echo $current_traffic | jq 'to_entries[] | .result = (.key|tostring)+ "=" + (.value|tostring)'  | jq '.result' -r | tr '\n' ',' | sed -e 's/,$/\n/g' )

# 新しいモデルのトラフィック割合を0にして、既存のモデルの割合を変更せずにデプロイ
gcloud ai endpoints deploy-model …… –traffic-split=0=0,${formatted_current_traffic}

さらなる注意点として、トラフィックが0%になった既存モデルのアンデプロイの自動化が挙げられます。トラフィックを完全に切り替えても、既存モデルは待機状態のままであるため、使用しているリソース分の料金が請求されます。 最初は定期的にモデルを監視して不要なモデルを手動でアンデプロイしていました。現在は使用していないモデルをアンデプロイするワークフローを作成してその作業を自動化しています。

追跡可能性

CIの中でモデルをデプロイしているため、CIのログ上には、いつどのコードを何にデプロイしているかの記録は残っています。しかし、それを目で辿っていくのは困難を極めます。

そこで、Artifact Registryへのpush時にgitのcommit hashをDockerのタグ情報として記録するようにしました。さらに、このcommit hashをVertex AI Model Registry上のモデルのエイリアスにも登録し、現在デプロイされているAPIから、作成したソースコードまでを簡単に辿れるようにしました。

一点気をつけることとして、Vertex AI Model Registryのエイリアスは英小文字始まりで指定する必要があります。ドキュメントに理由は見当たりませんでしたが、commit hashは数字で始まることもあるため、commit-というprefixを付けて回避しました。

Google Cloud コンソールからエイリアスを新規作成しようとした画面。英小文字指定が求められている。

デプロイ戦略

AI Labではモノレポを採用しており、APIソースコードは全て該当のレポジトリに入っています。レポジトリはmainブランチを1つだけ持っており、開発時はmainから派生しています。そしてmainブランチへのマージで開発環境にデプロイが走るようになっており、mainの最新を反映し続けるようにしています。 また、それ以外にも手元から特定の命名規則に従ったタグをpushすることでもデプロイ可能にしました。これによって開発時に手元からさっと開発環境にデプロイすることが可能です。複雑なブランチルールを持たず、タグpushだけで気軽にデプロイできるためデプロイのハードルはかなり低くなっていると感じています。

検証環境や本番環境へのデプロイも現在はタグpushを採用していますが、mainブランチ以外からのデプロイを避けるため今後はWorkflow dispatchでデプロイするように変更予定です。

ドキュメントツール

API基盤が作成されMLモデルをサービングできるようになりましたが、Vertex AI EndpointsはAPIの仕様に関しては面倒を見てくれません。入出力のスキーマはどのようなものなのか、どのカラムにどんな型のどんな値を期待しているのか、返してくれるのか、などは最小限のスキーマ情報を除いてほとんど未知の状態です。

そこで、APIの仕様を手で記述してドキュメントサーバ上に公開することにしました。幸いなことに、特定のGCSバケット上に静的ファイルを置くと簡単に社内ドキュメントサーバにデプロイできる仕組みがCADDi内にありました。我々のAPIドキュメントも同じように、GCS上に仕様を記述したHTMLファイルを置くことで社内公開することにしました。

APIの仕様の記述にはOpenAPIを使用しました。YAMLAPI情報を記述するのですが、APIの仕様記述に関して一般的なフォーマットであり、HTMLへのコンバートも簡単にできること、複数のYAMLファイルに分割して記述しても簡単にまとめることができること、が大きな選定理由です。OpenAPIに則ったYAMLファイルの管理ディレクトリの構造は以下のように決めました。 これらを swaggerを通して1つの大きなHTMLファイルへと変換しています。

APIの仕様にはVertex AI EndpointsのURL、入出力のスキーマ情報に加えて、入出力の各値の具体的な意味や実際の入出力例を記述するようにしました。というのも、単純なスキーマ情報では文字列型を期待していることはわかっても、実際にそこには画像のbase64エンコードテキストを入力することまでは分からないからです。

テストツール

ここまでで「APIを公開する」ということは可能になりましたが、すぐに「公開したAPIが本当に期待したものなのかのテストができていない」という問題に当たりました。 APIが本当に接続できる状態なのか、期待する出力が期待する時間内に返ってくるのか、等は不明なままです。具体的には、以下の内容を保証する必要が出てきました。

  • ドキュメントに書かれたAPIの型定義と、リクエスト・レスポンスの例に矛盾がない
  • ドキュメントに書かれたAPIのリクエスト例が、実際にリクエスト可能なものである
  • API基盤にデプロイしたモデルと、デプロイしようとした実験時のモデルの挙動が等しい
  • ある程度の高負荷時にも妥当なレスポンス時間で返ってくる

そこで、テストツールとして以下の内容をテストするツールを作成しました。

  • APIのドキュメントとして記述したOpenAPIが正しいフォーマットに沿っているかどうか
  • ドキュメントに書かれたリクエストの例を用いてAPIが疎通できるか
  • APIが、デプロイ前のMLモデルと同じ値を返せているか
  • APIに高負荷をかけてもリクエストが妥当な時間内に返ってくるか

OpenAPIのフォーマットテスト

OpenAPIとしての記述の正しさの確認にはswaggerを使用しました。しかし、このツールは型の中身についてチェックまではしてくれません。つまり、矛盾するような型定義を書いてもYAMLとして問題がなければチェックを通過してしまいます。そこで、さらに型定義のテストを行うようにしました。具体的には

  • リクエスト・レスポンスの型定義に矛盾がないか
  • リクエスト・レスポンスの具体例が型定義と一致しているか

型の定義はJSON schemaで記述されています。そこで、OpenAPIの型定義部分だけ抜き出して以下の内容をチェックするCLIを用意しました。チェックにはPythonモジュールの jsonschema を使用しました。

これにより、ドキュメントに書かれたAPIの型定義と、リクエスト・レスポンスの例の間に矛盾がないことが保証されました。

APIの疎通テスト

Vertex AI Endpointsにデプロイ後、サービングが確かに始まっていることを確認するため、サンプルリクエストを使用した疎通テストを追加しました。これによって、「デプロイしたと思っていたが実は途中で失敗したいた」という問題や、「ドキュメントのサンプルリクエストが実は間違っていた」という問題を消すことができました。

APIの性能テスト

MLエンジニアがモデルを作成した時の実行環境・ソースコード・ライブラリバージョンと、APIサーバのそれらを完全一致させるのは非常に難しいです。そのため、「作成したモデルとAPIの出力が一致しない」ということはしばしば起こり得ます。大切なのは、それらの誤差が許容できるレベルに十分小さいのか、もしくは何らかの問題により大きな誤差が発生しているのかを検知することです。

そこで、数十〜数百枚程度の小規模なデータセットに対して、モデル作成時にそれらの推論値をGCSに保存しておき、APIサーバがほぼ同じ出力を返すことを確認する性能テストを追加しました。APIサーバはローカル環境に立ち上げることで、デプロイ前にテストをすることが可能です。 これによって、API上でも期待する精度を出せることが保証できるようになりました。

APIの負荷テスト

これまでのテストで、APIが正しくデプロイされたことは保証されましたが、実際に使用に耐えうるレイテンシかどうかまでは確認できていません。そこで、APIドキュメントから使用したサンプルリクエストを使用した負荷テストを追加しました。 負荷テストのツールにはlocustを用いました。テスト実行後にHTMLでレポートが生成されるため、そのレポートを共有することでAPIが高負荷時にも許容できるレイテンシになっていることを保証しています。

API基盤のモニタリングツール

上記の複数のテストを経て、APIをデプロイすることができるようになりました。しかし、本当に大事なのはここからで、デプロイしたAPIを運用していく必要があります。具体的には、API基盤が問題を起こしていないかを監視し続ける必要があります。

監視にはCloud Monitoringを使用しました。Google Cloudで完結すること、Vertex AI Endpointsに対する基本的な監視がデフォルトで用意されていることが理由です。監視項目は以下の内容です。これらをダッシュボード上に用意し、業務時間中は常にサブディスプレイに表示し続けることで「Vertex AI Endpointsが正常に稼働している際にこれらの値がどのような傾向を示すか」というのを追い続けています。

監視をする中で、いくつか分かったことがあります。ここでは4つほどあげたいと思います。

GPUを使用しているAPIは、インスタンスのスケールアップが遅い

現在、CPU使用率が60%を上回ったAPIに関しては、インスタンス数を増やすように設定しています。しかし、実際にCPU使用率が60%を上回ってから、追加のインスタンスが確保されるまで、5-15分ほどがかかることがわかりました。図では5:43頃に閾値を超えていますが、実際にスケールアップしたのは5:52頃になっており、9分ほどかかっています。理由は分かっていませんが、TorchServeを含むDockerイメージのpullに時間がかかっているのかもしれません。

インスタンスのスケールアップ中はAPIのレスポンス速度が急激に低下する

CPU使用率が60%を超えてからインスタンスが増加するまで、APIのレスポンス速度が急激に低下します。デプロイしているAPIの多くの普段のレイテンシが大体0.1秒から1.0秒ですが、インスタンススケールアップ中は10-60秒ほどに低下します。また、この間に多くのリクエストに対して503エラーが返るようになります。 現時点では、5XX系のエラーに関しては利用者側にリトライ処理をお願いしています。

GPUのメモリ使用量はほとんど一定である

高負荷時も低負荷時も、GPUのメモリ使用量に大きな変化はなく、リクエストのボトルネックになることはなさそうでした。一方GPU使用率の方は負荷に応じた変動を示しています。

問題発生時にできることはほとんどない

Vertex AI Endpointsはフルマネージドサービスであるため、Vertex Modelとそれを動かすインスタンス条件だけを決めれば、あとはよしなにサービングしてくれます。逆に言えば、API基盤として問題が起きたとしてもやれることはほとんどなく、Google Cloud全体の障害かどうかを判断するかくらいです。もちろん精度の問題を検知した場合は、過去バージョンに巻き戻すのような対応がありますが、これも簡単に行うことができます。 とはいえ、まだ運用しはじめて3ヶ月ほどですが、基盤として問題が生じたことはなく安定したサービスだなと感じます。

Vertex EndpointのHuman readableな識別子がない

これは地味に辛い問題です。Vertex Endpointは識別子として数値列を使用しています。ユーザ側が指定することもできますし、特に指定しなければ適当な値を割り振られます。しかし、どちらの場合であっても数値列だけを見てどのAPIのものかを判断することは難しいです。そのため、APIには displayName という表示名を与えることができますが、これをCloud Monitoringから確認することができません。現時点ではダッシュボード上のテキストカードに、識別子とdisplayNameの対応を記述することで凌いでいます。

結果

半年前に3ヶ月ほどで作られたAPI基盤ですが、現在も開発や運用が続けられています。 まだ荒削りなとことはありますが、チームが増えた今でもシステム自体は容易に把握することができ、いくつものAPIがMLエンジニア主導でデプロイされていたり、新たにチームにジョインしたMLOpsメンバーが1週間で機能改善の本格的なPRを作成したりしてくれています。

一方、まだテストを自動化しきれていなかったり、スケールアップ時にAPI基盤が不安定になったりする問題があります。 今後は、テストやデプロイメントを改良しつつ、よりユーザに使ってもらいやすいAPI基盤を目指していこうと思います。

最後に

MLOPsチームではAPI基盤を皮切りに、ビジネスサイドに提供するためのバッチ推論基盤やデータ管理基盤にも取り組んでいます。今後は再学習基盤やデータ基盤など、より早く広く安全にAIの価値をビジネスに提供するための仕組みを作って行くので、一緒に作っていける仲間を募集しています!