Vertex AI Model Registry で機械学習モデルのバージョン管理をする

Vertex AI Model Registry とは

その名の通り、 ML モデルのバージョンを管理するための場所です。モデルのバージョン管理をするだけであれば GCS 等にモデルの情報を記載したテキストファイルを用意するなどの方法が考えられますが、 Vertex AI Model Registry を使うことの1番のメリットはやはり他の Vertex AI サービスとの親和性が高い点です。

例えば Model Registry に配置したモデルは以下のように利用できます。

1. テストデータに対するモデルの精度検証
2. Vertex AI Endpoint へのデプロイ
3. バッチ予測の実行

また BigQuery ML や AutoML で学習されたモデルも Model Registry で管理できるため、 GCP の何らかのサービスで学習した ML モデルは全て Model Registry に集約させることができます。

やりたいこと

この記事では「GCS で学習済みモデルのバージョン情報を管理している」というありがちな ML モデルの手動管理を Model Registry に移行する、というシチュエーションを考えます。

例えば GCS のバケットに対して以下のようなディレクトリ構成を取り、 latest ディレクトリにあるモデルを使うというのがモデル管理のありがちなパターンだと思っています。

.
└── bucket
    └── model
        ├── 20221201
        │   ├── info.txt
        │   └── model.pkl
        ├── 20221202
        │   ├── info.txt
        │   └── model.pkl
        ├── 20221203
        │   ├── info.txt
        │   └── model.pkl
        ├── 20221204
        │   ├── info.txt
        │   └── model.pkl
        └── latest
            ├── info.txt
            └── model.pkl  # ここを常に参照するようにしている

これをいい感じにシュッとさせることをゴールとします。

準備

必要なパッケージをインストールしておきます。

$ pip install google-cloud-aiplatform
$ pip list | grep aiplatform
google-cloud-aiplatform           1.18.2

モデルをアップロードする

GCS に配置したモデルを Model Registry にアップロードします。

from google.cloud import aiplatform

project = "your-project-id"
location = "your-location"
model_name = "awesome-model"
artifact_uri = "gs://bucket/path/to/model"
serving_container_image_uri = "us-docker.pkg.dev/vertex-ai/prediction/tf2-gpu.2-8:latest"
version_description = """
This is my first awesome model.
Trained on Vertex AI Workbench by xxx.
"""

aiplatform.init(project=project, location=location)
model = aiplatform.Model.upload(
    display_name=model_name,
    artifact_uri=artifact_uri,
    serving_container_image_uri=serving_container_image_uri,
    version_description=version_description,
)
print(model)

注意点

いきなり注意点です。ただモデルを Model Registry に登録するだけではありますが、サービングする時に使用するサービング用コンテナの URI を指定する必要があります。

上記の例では tensorflow のモデルを GPU インスタンス上でサービングすることを想定したイメージを指定しています。ここで指定するイメージは Vertex AI が提供しているビルド済みコンテナを指定するのが楽ですので、指定可能なイメージは以下から確認ください。

ビルド済みコンテナを使用する場合、artifact_uri に指定するパスの先に model.pkl や model.joblib が存在することを要件としている場合があります。
上記の例で使用している tensorflow のイメージの場合は saved_model.pb か saved_model.pbtxt が存在していることが要件となっています。
ビルド済みコンテナの要件はドキュメントを参照ください。

当然自前で用意したサービングコンテナを使用することもできますので、自分で Flask や Fast-API などのフレームワークを使用して自分でサービングコンテナを作る場合は以下のドキュメントを参照ください。

モデルの最新版をアップロードする

model_id の確認

最新版のモデルをデプロイするには、デプロイ済みモデルの model_id と呼ばれる値を確認する必要があります。

model_id はモデルの識別子のため、モデルのバージョン管理をする場合は同じ model_id に対して最新版をデプロイしていく必要があります。

model_id はデプロイ済みモデルのバージョンを UI から選択すると確認できます。

以下のように gcloud コマンドを使って確認することもできます。

$ gcloud ai models list --project=$PROJECT_ID --region=us-central1
Using endpoint [https://us-central1-aiplatform.googleapis.com/]
MODEL_ID: 8062923275644174336
DISPLAY_NAME: awesome-model

最新版モデルのデプロイ

先ほど確認した model_id を parent_model に指定して upload することで新しいバージョンとしてデプロイできます。

new_artifact_uri = "gs://bucket/path/to/new/model"
version_description = "Tuned hyper parameters as ..."

model = aiplatform.Model.upload(
    parent_model="8062923275644174336",
    artifact_uri=new_artifact_uri,
    serving_container_image_uri=serving_container_image_uri,
    version_description=version_description,
)

無事にデプロイが完了すると、新たにデプロイしたものがデフォルトモデルとして扱われます。ここまでの処理で、 UI 上では以下のような見え感になっています。見え感がちょっと寂しかったので Model Description に適当な文章を追加しました。

デプロイしたモデルをデフォルトとして扱わない場合

ここでは常に最新版を使う想定のためデプロイされたモデルをすぐに使うようにしていますが、 Model Registry に上げた後にテストデータでの検証などを行ったのちに UI 上でデプロイするような場合は is_default_version というフラグで制御できます。

以下のコマンドでは新しいバージョンをデプロイしていますがデフォルトバージョンとして扱わないようにしています。

new_artifact_uri = "gs://bucket/path/to/new/model"
version_description = "Please evaluate me using test data."

model = aiplatform.Model.upload(
    is_default_version=False,  # added
    parent_model="8062923275644174336",
    artifact_uri=new_artifact_uri,
    serving_container_image_uri=serving_container_image_uri,
    version_description=version_description,
)

デプロイされてもデフォルトのエイリアスが付与されていないことがわかります。

version alias でモデルを管理する

モデルのバージョンに対してエイリアスを貼ることで、例えば実験フェーズの experiment と運用フェーズの operation エイリアスをつけてモデルを識別することができます。

正直エイリアスを付けることによる恩恵を受けられる状況になったことがないので実際の運用では使用していませんが、ちゃんとエイリアスの運用設計をすれば受けられる恩恵もあるものと想像しています。

以下の要領でエイリアスを付与できます。

new_artifact_uri = "gs://bucket/path/to/new/model"
version_description = "Added version aliases"

model = aiplatform.Model.upload(
    parent_model="8062923275644174336",
    artifact_uri=new_artifact_uri,
    serving_container_image_uri=serving_container_image_uri,
    version_description=version_description,
    version_aliases=["experiment", "feature-set-v2"],
)

最新版のモデルを取得する

最後にここまでデプロイしてきたモデルを Model Registry から取得します。 model_id を指定して取得すると常にデフォルトバージョンのモデルが参照されます。

import pickle

model_id = "8062923275644174336"
_model = aiplatform.Model(model_name=str(model_id))
artifact_uri = _model.gca_resource.artifact_uri

with open(artifact_uri, "rb") as f:
    model = pickle.load(f)

注意点として Model Registry からはモデルそのものが取得できるわけではないため、モデルの本体が配置されている GCS からモデルファイルを読み込む必要があります。

また今回の場合は model_id を指定してデフォルトバージョンのみを取得していますが、aiplatform.Model.list メソッドで取得できるモデル群をフィルターして使用するモデルを選択する方法も考えられます。

今回のようなユースケースは想定されていないせいかフィルターの指定方法が詳しく書かれていなかったので、あまり推奨される方法ではないのかもしれません。

まとめ

• Model Registry を使うことでモデルのバージョン管理をいい感じにできる & 他の Vertex AI との親和性が高い
• 常に最新版のモデルを運用する場合は単一の model_id にデプロイし続けるだけ
• バージョン管理目的以外でも有用なケースはあるだろうが、今回はそこまで踏み込めなかった

Vertex AI との親和性の高さについて全く言及できなかったので、これはまた別な記事で取り上げたいと思います。