Zenn
❤️

Vertex AI Matching Engineを試す

2023/05/06に公開
GCP
Python
Vertex AI
#
matchingengine
tech

この記事について

GCPにはVertex AI Matching Engineという低レイテンシで近似最近傍探索を行なってくれるサービスがあります。レコメンドや検索に活用できそうなので気になっており、実際に使ってみました。

2023/05月の時点だと情報が少なく、いろんな事にハマったので記録として残しておきます。

構築する際の選択肢

ネットワーク関連

  • vpc peeringをして繋ぐ方法

  • public endpointという方法 (2023/05時点でプレビュー)

  • private service connectという方法 (2023/05時点でプレビュー)

の３つがあるらしい。

チュートリアルなどは主にvpc peeringをして繋ぐ方法を採用している。
しかし、VPCネットワークとかを構成する必要があり面倒臭いので、今回は、public endpointによる方法で進める。

インデックスやエンドポイントの作成方法

terraformによる対応はまだされていないみたい。

  • pythonのクライアントライブラリから作成する方法
  • curlやgcloudコマンドで作成する方法

前者は思ったより複雑なコードを書かなきゃいけないみたいだった。
後者は必要最小限の情報をjsonで定義してpostで送るため、わかりやすいと感じた。
そのため、今回は後者で実装する。

検索の方法

  • 自前でprotoファイル書いて検索
  • curlを使って検索
  • pythonのクライアントライブラリを使って検索

チュートリアルなどはprotoファイルを書いてコンパイルしてgRPC通信を行っている。
個人的には、検索の場合はクライアントライブラリを使ったほうが楽と判断したので、今回はこちらで進める。

構築の流れ

下記のプロセスで、検索を行うことができるようになる。

  1. データの準備
  2. インデックスの作成
  3. インデックスエンドポイント作成
  4. インデックスをエンドポイントにデプロイ
  5. 検索のためのclientを書く
  6. 検索クエリを実行

1. データの準備

100次元のランダムな値を持ったベクトルを100個ほど生成する。

import json
import random

# JSONファイルを開く
with open("index_data.json", "w") as f:
    # 改行区切りのJSONオブジェクトを含む文字列を生成する
    data_str = ""
    for i in range(100):
        obj = {
            "id": f"{i}",
            "embedding": [round(random.random(), 5) for _ in range(100)],
        }
        obj_str = json.dumps(obj)
        data_str += obj_str + "\n"

    # 文字列をファイルに書き込む
    f.write(data_str)

データ生成後は、GCSのバケット内にデータを置いておく。
なおマルチリージョンのバケットではなく、単一のリージョンでバケットを作成する必要がある。

2. インデックスの作成

インデックス用のjsonを定義する。
細かい意味などはドキュメントを参照してください。

index.json
{
    "contentsDeltaUri": "データが置いてあるGCSのURI",
    "config": {
      "dimensions": 100,
      "approximateNeighborsCount": 5,
      "shardSize": "SHARD_SIZE_SMALL",
      "algorithmConfig": { "treeAhConfig": {} }
    }
}

（認証周りのログインを済ませた上で）　コマンドを実行する。

gcloud ai indexes create \
  --metadata-file=index.json \
  --display-name=インデックスの名前 \
  --project=プロジェクト \
  --region=ロケーション

インデックスが準備されるまでそこそこ時間がかかる。

3. インデックスエンドポイント作成

エンドポイント用のjsonを作成する。
publicEndpointEnabledをtrueにすることによって、public endpointにすることを宣言している。

endpoint.json
{
   "display_name": "エンドポイントの名前", 
   "publicEndpointEnabled": "true" 
}

コマンドを実行する。

curl -X POST \
    -H "Authorization: Bearer $(gcloud auth print-access-token)" \
    -H "Content-Type: application/json; charset=utf-8" \
    -d @endpoint.json \
    "https://ロケーション-aiplatform.googleapis.com/v1/projects/プロジェクト/locations/ロケーション/indexEndpoints"

エンドポイントが作成されるのをまつ。

4. インデックスをエンドポイントにデプロイ

インデックスとエンドポイントの作成が完了したら、インデックスをエンドポイントにデプロイする

INDEX_ENDPOINT_IDとINDEX_IDは、先ほど作成した画面に表示されている。

DEPLOYED_INDEX_IDとDEPLOYED_INDEX_NAMEは、ここで文字列で新たに定義する必要がある。

（最初ID・DEPLOYEDとついているし、どこからか値を探してきて指定する必要があるのかと勘違いしてしまっていた。)

コマンドを実行する。

gcloud ai index-endpoints deploy-index INDEX_ENDPOINT_ID \
  --deployed-index-id=DEPLOYED_INDEX_ID \
  --display-name=DEPLOYED_INDEX_NAME \
  --index=INDEX_ID \
  --project=プロジェクト \
  --region=ロケーション

デプロイが完了するまで待つ。

public endpointで展開する場合、検索を行う際にpublicEndpointDomainNameの情報が必要になるので、下記のコマンドを叩いて確認しておく。

curl -H "Content-Type: application/json"  \
-H "Authorization: Bearer `gcloud auth print-access-token`"  \
https://ロケーション-aiplatform.googleapis.com/v1/projects/プロジェクト/locations/ロケーション/indexEndpoints/INDEX_ENDPOINT_ID

・・・・
  "publicEndpointDomainName": "hogehoge.ロケーション-fugafuga.vdb.vertexai.goog"
}

5. 検索のためのclientを書く

  • sa_file_path

GCPでVertex AI Userのroleを持つサービスアカウントを作成してサービスアカウントキーをダウンロードしておいたパス

  • api_endpoint

先ほどのpublicEndpointDomainNameを指定

from google.cloud import aiplatform_v1beta1
from google.oauth2 import service_account


def get_client():
    # The AI Platform services require regional API endpoints.
    scopes = ["https://www.googleapis.com/auth/cloud-platform"]

    # create a service account with `Vertex AI User` role granted in IAM page.
    # download the service account key https://developers.google.com/identity/protocols/oauth2/service-account#authorizingrequests
    sa_file_path = "サービスアカウントのクレデンシャルファイル"

    credentials = service_account.Credentials.from_service_account_file(
        sa_file_path, scopes=scopes
    )
    client_options = {"api_endpoint": "先ほど確認したpublicEndpointDomainName"}

    vertex_ai_client = aiplatform_v1beta1.MatchServiceClient(
        credentials=credentials,
        client_options=client_options,
    )

    return vertex_ai_client


client = get_client()

6. 検索クエリを実行

import random

request = aiplatform_v1beta1.FindNeighborsRequest(
    index_endpoint="projects/プロジェクト/locations/ロケーション/indexEndpoints/INDEX_ENDPOINT_ID",
    deployed_index_id="DEPLOYED_INDEX_ID",
)
dp1 = aiplatform_v1beta1.IndexDatapoint(
    datapoint_id="0",
    feature_vector=[round(random.random(), 5) for _ in range(100)],
)
query = aiplatform_v1beta1.FindNeighborsRequest.Query(
    datapoint=dp1,
)
request.queries.append(query)

response = client.find_neighbors(request)

検索結果が返ってくる

おまけになるが、上記のレスポンスは、protobuf形式らしく、pythonで扱いづらい。
from google.protobuf.json_format import MessageToJsonなどでパースできるらしいが、どうにもうまくいかない。そのため違うライブラリを使ってパースしている。今回の主目的ではないのでスルー。

import proto

parsed_response = [proto.Message.to_dict(search) for search in response.nearest_neighbors]

お片付け

デプロイされたインデックスを解除してから、インデックスやエンドポイントを削除できる

デプロイされたインデックスの解除

gcloud ai index-endpoints undeploy-index INDEX_ENDPOINT_ID \
  --deployed-index-id=DEPLOYED_INDEX_ID \
  --project=プロジェクト \
  --region=ロケーション

エンドポイントの削除

gcloud ai index-endpoints delete INDEX_ENDPOINT_ID \
--project=プロジェクト \
--region=ロケーション

インデックスの削除

gcloud ai indexes delete INDEX_ID \
--project=プロジェクト \
--region=ロケーション

あとはいらないGCSバケットやサービスアカウントなど削除！

所感

web上に情報が少なく、ドキュメントもよくわらかないところが多いので、ハマった時に情報がなく解決が大変で非常に辛い。
そこさえ乗り越えられるなら、中で使われているScaNNが定量的には性能が良いらしいし、今回紹介していない欲しい機能も揃えている。そのため、自前で作成するのではなく、Matching engineを使うというのも十分選択肢になりうると感じた。

Discussion