Vertex AI Matching Engineの検索速度を試す

バックエンドエンジニアの@nasaです。Google Cloudの近似最近傍探索サービスのVertex AI Matching Engineを試したのでその話を書こうと思います。

Vertex AI Matching Engine(以下Matching Engine)の自体の説明はしないつもりなので詳しく知りたい人は公式ブログを見てみてください〜

本記事ではMatching Engineのindex作成速度、検索速度について書こうと思います。

TL;DR

768次元ベクトル10万件でのindex作成は50分ほど
デプロイは5分ほど。ただしindex作成からしばらくたってデプロイすると50分ほど(おそらくindexを再構築している)
1万件取得する際の検索速度は12msほど

検索までのオーバーヘッドはあるものの検索自体は爆速で終わりますね。取得件数が多くても問題なく捌けているので大規模な検索にも使えそうです。

Matching Engineでベクトル検索するまでの流れと速度

軽くMatching Engineでベクトル検索するまでの流れをまとめようと思います。その後、どの工程にどれだけ時間がかかるのかまとめます。

次の5ステップを行なうことでベクトル検索を行うことが出来ます。

ベクトルをGoogle Cloud Storageにupload
Google Cloud Storageのデータを元にMatching Engineのindexを作成
indexをデプロイするためのindex endpointを作成
indexをindex endpointにデプロイ(ここまでやるとgRPCエンドポイントが作られる)
gRPCエンドポイントに検索リクエストを送信

1~4は一度行えば十分ですが、データ更新があった際にはindexを再計算する必要があります。この際、indexがデプロイ済みであればindexの更新後に自動でデプロイされます。

1のCloud Storageにアップロードする部分はデータ量依存なのとMatching Engineとは関係ないところなので省略します。また3のindex endpointの作成は一瞬で終わるので省略します。

ここでちょっとindex endpointについて説明します。

index endpointは生成したindexを検索するために作るリソースです。(indexを検索可能にする箱のようなイメージ)。indexをindex endpointにデプロイすることで初めてベクトル検索ができるようになります。1つのindex endpointに複数のindexをデプロイすることが出来るので1つだけ作っておけば良いものだと思っています。

それでは、2,4,5のindex作成、indexデプロイ、検索の3つのを見ていきましょう。今回計測するものは、上記のステップに沿って書くと次のようになります。

ベクトルをGoogle Cloud Storageにupload(計測しない)
indexを作成 (計測)
indexをデプロイするためのindex endpointを作成 (計測しない)
indexをindex endpointにデプロイ (計測)
gRPCエンドポイントに検索リクエストを送信 (計測)

index作成時間

次のような条件で計測を行いました。

ベクトルは768次元
ベクトル数は10万件
これらのベクトルはrand関数で適当に生成したものを使用する(0~1の範囲の値を生成)
Google Cloud Storageにはjson形式で保存している

indexの作成はpythonで行いました。Google Cloudが提供しているpython clientを使いindex作成apiを叩いています。

index作成時のパラメータはサンプルコードで使用されているものを拝借しました。

各値を意味はドキュメントを見てみてください https://cloud.google.com/vertex-ai/docs/matching-engine/using-matching-engine?hl=ja#index-metadata-file

from google.cloud import aiplatform_v1beta1
from google.protobuf import struct_pb2

GCS_URI  = "<ベクトルがアップロードされているGCSのリンク>"
PARENT = "projects/{}/locations/{}".format("<GCPのプロジェクト名>" ,"<GCPリージョン>")
ENDPOINT ="{}-aiplatform.googleapis.com".format("<GCPリージョン>")


treeAhConfig = struct_pb2.Struct(
    fields={
        "leafNodeEmbeddingCount": struct_pb2.Value(number_value=500),
        "leafNodesToSearchPercent": struct_pb2.Value(number_value=7),
    }
)

algorithmConfig = struct_pb2.Struct(
    fields={"treeAhConfig": struct_pb2.Value(struct_value=treeAhConfig)}
)

config = struct_pb2.Struct(
    fields={
        "dimensions": struct_pb2.Value(number_value=768),
        "approximateNeighborsCount": struct_pb2.Value(number_value=150),
        "distanceMeasureType": struct_pb2.Value(string_value="DOT_PRODUCT_DISTANCE"),
        "algorithmConfig": struct_pb2.Value(struct_value=algorithmConfig),
    }
)

metadata = struct_pb2.Struct(
    fields={
        "config": struct_pb2.Value(struct_value=config),
        "contentsDeltaUri": struct_pb2.Value(string_value=GCS_URI),
    }
)

ann_index = {
    "display_name": "dummy-index",
    "description": "適当にindexを作るぜ",
    "metadata": struct_pb2.Value(struct_value=metadata),
}

index_client = aiplatform_v1beta1.IndexServiceClient(
    client_options=dict(api_endpoint=ENDPOINT)
)

index_client.create_index(parent=PARENT, index=ann_index)
while True:
    if ann_index.done():
        break
    logx.debug("気長にお待ちを ><")
    time.sleep(60)

上記の前提でindexを作成した場合は51分かかりました。

(データ量によってどこまで変わるのか、パラメーターによって変わるのかは見れていません。早くなると嬉しいのでちょっと試してみたいですね。)

indexデプロイ時間

indexは先ほど作成したものを使用します。これのデプロイにどれくらい掛かるのでしょう。

同じくpythonスクリプトを使ってデプロイしてみましょう

from google.cloud import aiplatform_v1beta1

ENDPOINT ="{}-aiplatform.googleapis.com".format("<GCPリージョン>")
index_resource_name = "<index名>"
DEPLOYED_INDEX_ID = "<デプロイしたindexのid>" 

index_endpoint_client = aiplatform_v1beta1.IndexEndpointServiceClient(
    client_options=dict(api_endpoint=ENDPOINT)
)

deploy_ann_index = {
    "id": DEPLOYED_INDEX_ID,
    "display_name": "",
    "index": index_resource_name,
}

r = index_endpoint_client.deploy_index(
    index_endpoint=endpoint_name, deployed_index=deploy_ann_index
)

while True:
    if r.done():
        break
    print("ちょっとまってね")
    time.sleep(60)

index作成直後は5分程で完了しました。

ただ、index作成から1日経過したときの実行時間は50分ほどでindex作成とほぼ同等の時間になりました。推測ですが、作成したindexをデプロイしないままほっておくとindexが保持されずデプロイ時に再構築されているのではないかと考えています。

検索速度

メインディッシュの検索速度ですね。

公式ブログには10ms以下のレスポンスタイムだと言っていますが果たして！

95〜98 % の再現率と 10ms 以下のレスポンスタイムを提供します。

ref: https://cloud.google.com/blog/ja/topics/developers-practitioners/find-anything-blazingly-fast-googles-vector-search-technology

indexは先程作ったものを使用しました。

計測はVertex AI Workbench上でpythonを動かして行いました。Workbenchは次のようになっています。

環境: TensorFlow Enterprise 2.8 (with LTS and Intel® MKL-DNN/MKL)
環境バージョン: M90
マシンタイプ: n1-standard-4 (4 vCPUs, 15 GB RAM)

コードは下記のものを使用しました。検索クエリのベクトルはランダムに生成したものを使用しています。また、レスポンスは1万件に設定しています。

from google.cloud import aiplatform_v1beta1
import match_service_pb2      # サンプルコードにあるprotoから生成
import match_service_pb2_grpc # サンプルコードにあるprotoから生成

DEPLOYED_INDEX_ID = "<デプロイしたindexのid>" 
DEPLOYED_INDEX_SERVER_IP = "<gRPCエンドポイントのアドレス>"

channel = grpc.insecure_channel("{}:10000".format(DEPLOYED_INDEX_SERVER_IP))
stub = match_service_pb2_grpc.MatchServiceStub(channel)

queries = [] 
for _ in range(100):
    query = [random.random() for _ in range(768)]
    queries.append(query)

hits = 0
all_time = 0
for query in queries:
    request = match_service_pb2.MatchRequest()
    request.deployed_index_id = DEPLOYED_INDEX_ID
    request.num_neighbors = 10000
    
    for val in query:
        request.float_val.append(val)
    time_sta = time.perf_counter()
    response = stub.Match(request)
    time_end = time.perf_counter()
    tim = time_end- time_sta
    hits += len(response.neighbor)
    all_time += tim
    

print("avg time", all_time / len(queries), "s")
=> avg time 0.012501707095652818 s

結果は1リクエストあたり12msでした。おおー。爆速なのでは。

(結果を紛失しているので正確な値ではないですが、取得結果が100, 1000のときも同じくらいのレスポンスタイムでした)

少し前にElasticsearchを試したときは1万件取得に400msほどかかっていたのでそれと比べるとかなり早いですね。

Elasticsearch v8に導入されたkNN search APIの検索速度を試す | Wantedly Engineer Blog

バックエンドエンジニアの@nasa です。最近Elasticsearchのversion8から追加された近似最近傍探索(以下ANN)を試してみたのでその話をしようと思います。 Elasticsearch version 8から追加された knn search API にリクエストを送信して速度検証を行いましたANN自体については取り扱わないので、知りたい人はElasticsearchの公式ドキュメントを参照してください。 ANNをパッと試したときの記録があるので、実際に動かしてみたい人はこちらを参照し

https://www.wantedly.com/companies/wantedly/post_articles/386300