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

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

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

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

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

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

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

```shell
.
└── 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  # ここを常に参照するようにしている
```

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/265024/4fba8eb8-dc09-a48b-0000-456a7057386e.png)


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

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

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

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

```python
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 が提供しているビルド済みコンテナを指定するのが楽ですので、指定可能なイメージは以下から確認ください。

https://cloud.google.com/vertex-ai/docs/predictions/pre-built-containers

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

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

https://cloud.google.com/vertex-ai/docs/predictions/custom-container-requirements

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

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

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

![スクリーンショット 2022-12-14 0.20.57.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/265024/a784da2c-95bd-ae92-6187-03e3b083acb8.png)

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

```shell
$ 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 することで新しいバージョンとしてデプロイできます。

```python
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 に適当な文章を追加しました。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/265024/46554ed4-25c6-4cc6-6674-129b3c64ad36.png)

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

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

```python
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,
)
```

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

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/265024/11d8bf06-6c97-5fb4-5818-90dd32776a02.png)

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

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

https://cloud.google.com/vertex-ai/docs/model-registry/model-alias

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

```python
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"],
)
```

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/265024/033eb9c3-3d6d-d1b4-f7e4-99e9f8aa9eeb.png)


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

```python
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 との親和性の高さについて全く言及できなかったので、これはまた別な記事で取り上げたいと思います。