🍭

「Vertex AIで学習しないとModel Registryは使えない」という誤解を解く。自前モデルを管理するだけの「いいとこ取り」活用術

に公開

はじめに

こんにちは! AI技術開発室のYTです。機械学習モデルを本番環境で運用する際、モデルのバージョン管理は非常に重要です。どのバージョンのモデルがいつデプロイされ、どのような性能だったかを追跡できなければ、再現性の確保や問題発生時の原因究明が困難になります。

Google Cloudには、こうしたモデル管理の課題を解決するためのサービスとしてVertex AI Model Registryがあります。しかし、「Vertex AIで学習や予測を行っていないと使えないのでは?」と思っている方もいらっしゃるかもしれません。

実は、Vertex AIのパイプラインなどを使っていなくても、Vertex AI Model Registryの恩恵を受けることは可能です。今回は、Cloud Run上で動作するCatBoostモデルを例に、学習・予測をVertex AI上で行っていない機械学習モデルをVertex AI Model Registryで管理する方法をご紹介します。

Vertex AI Model Registryとは

Vertex AI Model Registryは、機械学習モデルを一元的に管理、バージョニングするためのリポジトリサービスです。モデルのライフサイクル全体を管理し、MLOpsのプラクティスを実践する上で中心的な役割を果たします。

Google Cloud Storageとの連携

Vertex AI Model Registryは、Google Cloud Storage (GCS) に保存されたモデルアーティファクトを登録することで機能します。特別な形式は必要なく、モデルファイル一式が格納されたGCSディレクトリを指定するだけで登録できます。

例えば、以下のようなディレクトリ構成でモデルファイルがGCSに保存されていれば、gs://your-bucket/models/catboost-regressor/20240101120000/ を指定することで、このバージョンのモデルをModel Registryに登録できます。

your-bucket/
└── models/
    └── catboost-regressor/
        ├── 20240101120000/
        │   ├── model.cbm
        │   └── metadata.json
        └── 20240102150000/
            ├── model.cbm
            └── metadata.json

実装の全体像

今回は、既存のCloud RunアプリケーションにVertex AI Model Registryを組み込む改修を行いました。

変更前の処理フロー

  1. 学習: Cloud Runでモデルを学習し、GCSの特定のパスにモデルファイルを保存する。
  2. 予測: 予測リクエストがあると、Cloud RunはGCS上の「最新」のモデルファイルをダウンロードして予測を実行する。

変更後の処理フロー

  1. 学習: 従来通りCloud Runでモデルを学習し、GCSにモデルファイルを保存する。
  2. モデル登録: 精度検証などをパスしたモデルを、Vertex AI Model Registryに新しいバージョンとして登録する。この際、productionなどのバージョンエイリアスを付与する。
  3. 予測: 予測リクエストがあると、Cloud RunはVertex AI Model Registryに問い合わせ、productionエイリアスを持つモデルバージョンのGCS URIを取得し、そのURIからモデルファイルをダウンロードして予測を実行する。

※ 今回扱うモデルは軽量であり、利用用途も社内向けサービスであるため、実装の複雑さを排除してリクエストごとに最新モデルを取得する(キャッシュを持たない)構成としています。これにより、エイリアスを切り替えた瞬間に全リクエストへ新モデルが適用されます。

変更前後の処理フロー図

Vertex AI Model Registryへモデルを登録する

精度検証をクリアしたモデルをModel Registryに登録する処理を実装します。Google CloudのPythonクライアントライブラリ google-cloud-aiplatform を使用します。

以下のコードは、指定されたGCSのパスにあるモデルを、指定したディスプレイ名でModel Registryにアップロード(登録)する関数の例です。

register_model.py
from google.cloud import aiplatform
import datetime

def register_and_promote_model_registry(
    project_id: str,
    location: str,
    model_display_name: str,
    gcs_artifact_directory: str,
    validation_metrics: dict,
    alias: str = "default",
):
    aiplatform.init(project=project_id, location=location)

    existing_models = aiplatform.Model.list(
        filter=f'display_name="{model_display_name}"',
        order_by="create_time desc"
    )
    
    parent_model_resource = None
    if len(existing_models) > 0:
        parent_model_resource = existing_models[0].resource_name

    timestamp = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    description = f"Validated at {timestamp}. Metrics: {validation_metrics}"

    # 推論はCloud Runで行うため、Vertex AI上ではダミーのコンテナイメージを登録する
    DUMMY_SERVING_CONTAINER = "gcr.io/google-containers/pause:latest"

    model = aiplatform.Model.upload(
        display_name=model_display_name,
        artifact_uri=gcs_artifact_directory,
        serving_container_image_uri=DUMMY_SERVING_CONTAINER,
        parent_model=parent_model_resource,
        description=description,
        version_aliases=[alias],
    )

    model.wait()

    print("\n✅ Model registration and promotion completed!")
    print(f"  Model Name: {model.display_name}")
    print(f"  Version ID: {model.version_id}")
    print(f"  Alias '{alias}' has been set to version {model.version_id}.")

    return model

Model Registryからモデルを取得して予測する

次に、予測時にModel Registryから本番用のモデルURIを取得する処理です。モデルのディスプレイ名とバージョンエイリアス(例: production)を使って、目的のモデルバージョンを特定します。

get_model_uri.py
from google.cloud import aiplatform
from google.api_core import exceptions
from typing import Optional

def get_model_uri_by_display_name_and_alias(
    project_id: str, 
    location: str, 
    model_display_name: str, 
    version_alias: str
) -> str:
    aiplatform.init(project=project_id, location=location)

    print(f"Searching for model with display_name: '{model_display_name}'...")

    models = aiplatform.Model.list(
        filter=f'display_name="{model_display_name}"',
        order_by="create_time desc" 
    )

    if not models:
        raise ValueError(f"Model with display_name '{model_display_name}' not found.")

    model = models[0]
    if len(models) > 1:
        print(f"Warning: Found {len(models)} models. Using the latest created one: {model.resource_name}")
    
    print(f"Found model group: {model.resource_name}")
    print(f"Getting version with alias: '{version_alias}'...")
    
    # モデルIDにエイリアスを付与してリソース名を構築
    model_version_resource_name = f"{model.resource_name}@{version_alias}"

    try:
        model_version = aiplatform.Model(model_name=model_version_resource_name)
        
        print(f"Found version: {model_version.version_id}")

        if not model_version.uri:
            raise ValueError(
                f"Artifact URI is empty for model version {model_version.version_id} (alias: '{version_alias}')."
            )

        return model_version.uri

    except exceptions.NotFound:
        raise ValueError(
            f"Alias '{version_alias}' not found in model '{model_display_name}' ({model.resource_name})."
        )
    except Exception as e:
        print(f"An unexpected error occurred: {e}")
        raise

この関数を使えば、予測サーバーは「productionエイリアスを持つモデル」という抽象的な指定でモデルを取得でき、具体的なGCSパスを意識する必要がなくなります。

この実装のメリットと注意点

メリット

  • モデルのバージョニングと再現性向上: どのモデルがいつ登録されたかが明確になり、過去のバージョンへの切り戻しも容易になります。
  • デプロイプロセスの簡素化: 予測サーバーは「productionのモデルを使う」とだけ知っていればよく、新しいモデルを本番投入する際はModel Registry上でエイリアスを付け替えるだけで完了します。アプリケーションの再デプロイは不要です。
  • 即時反映とステートレス性: リクエストごとにモデルを参照するため、Registry側でエイリアスを切り替えると、インスタンスの再起動を待つことなく、即座にすべてのリクエストで新しいモデルが使用されます。キャッシュ破棄のタイミングなどを考慮する必要がなく、非常にシンプルな構成を維持できます。
  • コストパフォーマンスの最適化: Vertex AI Endpoint(オンライン予測)を使用する場合、ノードが常時稼働するため待機コストが発生しますが、この構成(Model Registry + Cloud Run)であれば、リクエストがない時間はコストがゼロになります(Scale to Zero)。社内ツールのような「利用頻度は高くないが、いつでも使える状態にしたい」ケースでは、圧倒的なコストメリットがあります。

注意点

  • 権限管理: Cloud Runのサービスアカウントに、モデルファイルが格納されているGCSバケットへのアクセス権限だけでなく、Vertex AI Model Registryへのアクセス権限を付与する必要があります。
  • エイリアスの管理: どのエイリアスがどの環境(開発、ステージング、本番など)に対応するかのルールをチーム内で明確にしておくことが重要です。
  • Cloud StorageとRegistryの整合性: Vertex AI Model RegistryはGCS上のファイルを参照しているだけです。もし、誤ってGCS上の元ファイルを削除したり移動したりしてしまうと、Registry上は「有効」に見えても、予測時に FileNotFoundError が発生します。GCSバケットには削除防止のポリシーを設定するか、運用ルールで「Registry登録済みのファイルは触らない」ことを徹底する必要があります。

まとめ

今回は、Vertex AIの学習・予測機能を使わない既存のMLシステムに、Vertex AI Model Registryを導入する方法をご紹介しました。

Model Registryを活用することで、GCS上にファイルを置くだけの管理方法から脱却し、より堅牢で管理しやすいモデル運用体制を構築できます。Vertex AIの機能を部分的に採用するだけでも、MLOpsの成熟度を一段階引き上げることが可能です。ぜひお試しください。

We are hiring!

AI技術開発室では、一緒に働く仲間を募集しています。詳しくは以下をご覧ください。
https://hrmos.co/pages/carenet5800/jobs/1826582723293220966
https://hrmos.co/pages/carenet5800/jobs/0000020

CareNet Engineers

Discussion