背景

この記事では、Milvusがどのように問い合わせタスクをスケジューリングするかについて説明します。また、Milvusスケジューリングの問題点、解決策、今後の方向性についても説明します。

背景

大規模ベクトル検索エンジンにおけるデータ管理から、ベクトル類似性検索は高次元空間における2つのベクトル間の距離によって実行されることが知られている。ベクトル検索の目的は、ターゲットベクトルに最も近いK個のベクトルを見つけることである。

ユークリッド距離のように、ベクトル距離を測定する多くの方法があります：

1-euclidean-distance.png

ここでxとyは2つのベクトルであり、nはベクトルの次元である。

データセット内のK個の最近接ベクトルを見つけるには、ターゲットベクトルと検索対象のデータセット内のすべてのベクトルとの間のユークリッド距離を計算する必要がある。そして、K個の最近接ベクトルを求めるために、ベクトルを距離でソートする。計算量はデータセットのサイズに比例する。データセットが大きければ大きいほど、クエリに必要な計算量は多くなる。グラフ処理に特化したGPUは、必要な計算能力を提供するために多くのコアを持っている。そのため、Milvusの実装ではマルチGPUのサポートも考慮されている。

基本概念

データブロック（TableFile）

Milvusでは、大規模データ検索への対応を強化するため、データ保存の最適化を行った。Milvusはテーブル内のデータをサイズごとに複数のデータブロックに分割する。ベクトル検索では、Milvusは各データブロック内のベクトルを検索し、その結果をマージする。1回のベクトル検索操作は、N個の独立したベクトル検索操作(Nはデータブロック数)とN-1個の結果のマージ操作から構成される。

タスクキュー（TaskTable）

各Resourceは、そのResourceに属するタスクを記録するタスクアレイを持つ。各タスクには、開始、ロード、ロード済み、実行、実行済みなどの状態があります。コンピューティングデバイスのローダーとエクゼキュータは同じタスクキューを共有する。

クエリスケジューリング

2-query-scheduling.png

1. Milvusサーバーが起動すると、Milvusはserver_config.yaml 設定ファイルのgpu_resource_config パラメータを介して対応するGpuResourceを起動します。DiskResource と CpuResource はまだserver_config.yaml で編集できません。GpuResourceはsearch_resources とbuild_index_resources の組み合わせであり、以下の例では{gpu0, gpu1} と呼ぶ：

3-サンプルコード.png

3-example.png

2. milvusはリクエストを受け取る。テーブルメタデータは外部データベースに格納され、シングルホストではSQLiteまたはMySQl、分散ホストではMySQLが使用される。検索要求を受け取ると、Milvusはテーブルが存在し、ディメンションが整合しているかどうかを検証します。その後、MilvusはテーブルのTableFileリストを読み込みます。

4-milvus-reads-tablefile-list.png

3. milvusはSearchTaskを作成する。各TableFileの計算は独立して行われるため、Milvusは各TableFileに対してSearchTaskを作成する。タスクスケジューリングの基本単位であるSearchTaskには、ターゲットベクトル、検索パラメータ、TableFileのファイル名が含まれる。

5-table-file-list-task-creator.png

4. Milvusは計算デバイスを選択する。SearchTaskが計算を行うデバイスは、各デバイスの推定完了時間に依存する。完了予想時刻とは、現在時刻から計算が完了すると予想される時刻までの推定間隔を示す。

たとえば、SearchTaskのデータ・ブロックがCPUメモリにロードされたとき、次のSearchTaskはCPUの計算タスク・キューで待機しており、GPUの計算タスク・キューはアイドル状態です。CPUの推定完了時間は、前のSearchTaskと現在のSearchTaskの推定時間コストの合計に等しい。GPUの推定完了時間は、データ・ブロックがGPUにロードされる時間と、現在のSearchTaskの推定時間コストの合計に等しい。リソース内のSearchTaskの推定完了時間は、リソース内のすべてのSearchTaskの平均実行時間に等しくなります。そして、Milvusは、推定完了時間が最も短いデバイスを選択し、そのデバイスにSearchTaskを割り当てる。

ここでは、GPU1の推定完了時間の方が短いと仮定します。

6-GPU1-shorter-estimated-completion-time.png

5. MilvusはSearchTaskをDiskResourceのタスクキューに追加します。

6. MilvusはSearchTaskをCpuResourceのタスクキューに移動。CpuResourceのロードスレッドがタスクキューから各タスクを順次ロードする。CpuResourceが対応するデータブロックをCPUメモリに読み込む。

7. MilvusはSearchTaskをGpuResourceに移す。GpuResourceのロードスレッドがCPUメモリからGPUメモリにデータをコピーする。GpuResourceが対応するデータブロックをGPUメモリに読み込む。

8. milvusはGpuResourceでSearchTaskを実行します。SearchTaskの結果は比較的小さいため、結果は直接CPUメモリに返されます。

7-scheduler.png

9. MilvusはSearchTaskの結果を検索結果全体にマージする。

8-milvus-merges-searchtast-result.png

すべてのSearchTaskが終了すると、Milvusは検索結果全体をクライアントに返します。

インデックス作成

インデックス構築は基本的にマージ処理を除いた検索処理と同じです。これについては詳しく触れません。

パフォーマンスの最適化

キャッシュ

前述したように、データブロックは計算前にCPUメモリやGPUメモリなどの対応するストレージデバイスにロードされる必要があります。データロードの繰り返しを避けるため、MilvusはLRU（Least Recently Used）キャッシュを導入しています。キャッシュがいっぱいになると、新しいデータ・ブロックが古いデータ・ブロックを押しのけます。キャッシュのサイズは、現在のメモリサイズに基づいて設定ファイルによってカスタマイズすることができます。データのロード時間を効果的に節約し、検索パフォーマンスを向上させるには、検索データを保存するための大きなキャッシュを推奨します。

データの読み込みと計算の重複

キャッシュは検索パフォーマンスを向上させるニーズを満たすことはできません。メモリが不足したり、データセットのサイズが大きすぎる場合、データを再ロードする必要がある。データロードが検索性能に与える影響を減らす必要がある。ディスクからCPUメモリへ、あるいはCPUメモリからGPUメモリへのデータロードは、IOオペレーションに属し、プロセッサによる計算作業はほとんど必要ありません。そこで、データロードと計算を並列に実行することで、リソースの有効利用を図ります。

データブロックの計算を3ステージ（ディスクからCPUメモリへのロード、CPU計算、結果マージ）または4ステージ（ディスクからCPUメモリへのロード、CPUメモリからGPUメモリへのロード、GPU計算と結果検索、結果マージ）に分割します。3ステージの計算を例にとると、3つのステージを担当する3つのスレッドを起動し、命令パイプラインとして機能させることができます。結果セットはほとんど小さいので、結果のマージにはそれほど時間がかからない。場合によっては、データロードと計算のオーバーラップにより、検索時間を1/2に短縮できる。

9-sequential-overlapping-load-milvus.png

問題点と解決策

伝送速度の違い

以前、MilvusはマルチGPUタスクのスケジューリングにラウンドロビン戦略を使用していました。この戦略は我々の4GPUサーバーでは完璧に機能し、検索性能は4倍向上した。しかし、2GPUのホストでは2倍の性能向上には至りませんでした。私たちはいくつかの実験を行い、あるGPUのデータコピー速度が11GB/秒であることを発見しました。しかし、別のGPUでは3GB/秒でした。メインボードのドキュメントを参照した結果、メインボードがPCIe x16経由で1つのGPUに、PCIe x4経由でもう1つのGPUに接続されていることを確認した。つまり、これらのGPUはコピー速度が異なる。その後、各SearchTaskに最適なデバイスを測定するため、コピー時間を追加した。

今後の課題

複雑化するハードウェア環境

実際の環境では、ハードウェア環境がより複雑になる可能性がある。複数のCPU、NUMAアーキテクチャのメモリ、NVLink、NVSwitchを持つハードウェア環境では、CPU/GPU間の通信が最適化の機会を多くもたらします。

クエリの最適化

実験中に、パフォーマンス向上の機会をいくつか発見しました。例えば、サーバーが同じテーブルに対する複数のクエリを受信した場合、条件によってはクエリをマージすることができます。データの局所性を利用することで、パフォーマンスを向上させることができる。現在、シングルホスト、マルチGPUのシナリオにおいて、クエリがどのようにスケジューリングされ、実行されるかは既に分かっています。今後もMilvusの内部メカニズムを紹介していく予定である。