Milvus 查詢任務調度原理


Milvus 查詢任務調度原理

本文主要闡述在單機多卡的場景下,Milvus 是如何調度查詢任務的。此外,我們還將討論在調度實現過程中遇到的問題,解決方案和未來的發展方向。


| 背景

前文 大規模向量檢索場景下的數據管理(上篇)提到過,向量的相似性搜索實際上是通過兩個向量在高維空間的距離來判斷的,向量搜索就是在高維空間中找到與目標向量距離最接近的 K 個向量。

衡量向量相似性的距離有很多種,以最典型的歐式距離為例,其計算公式為:

Milvus 查詢任務調度原理


其中 x 和 y 是兩個向量,n 是向量的維度。

為了能找到集合中與目標向量最相似的 K 個向量,通常需要計算目標向量與集合中所有向量的歐式距離,然後根據距離對向量進行排序,獲得距離最小的 K 個向量,即得到與目標向量最相似的 K 個向量。向量搜索的計算量與集合規模成正比,搜索的集合規模越大,所需要的計算量就越大。而 GPU 作為圖形處理專用處理器,擁有大量的計算核心,正好可以提供計算所需要的大量算力,所以 Milvus 在實現的時候把多 GPU 異構計算的支持也考慮了進去。


| 基本概念

數據塊(TableFile)

為了更好地支持大規模的數據檢索,我們在 Milvus 的數據存儲上做了一些工作。對於一個 Table 的數據,我們會在插入時按照大小進行分割,得到多個數據塊。在進行向量搜索的時候,我們會在每一個數據塊中進行目標向量的搜索,最後把每個數據塊中獲得的結果歸併到一起得到最終結果。所以一次向量搜索的計算過程由 N 次獨立的向量搜索( N 為數據塊個數)和 N-1 次結果歸併組成。

計算設備(Resource)

為了能夠在不同硬件場景以配置的方式啟動計算資源,我們對計算設備進行了抽象。定義了包括 DiskResource、CpuResource 和 GpuResource。一個 CpuResource 代表一顆 CPU,一個 GpuResource 代表一張 GPU 卡。Resource 啟動好後就可以接受搜索任務了,每個 Resource 內部包含一個 Loader 和一個 Executor,Loader 負責將任務隊列中的任務數據加載到當前設備上,Executor 負責執行已加載任務的搜索。

任務隊列(TaskTable)

每一個 Resource 上都有一個任務隊列,任務隊列中會記錄屬於這個 Resource 的任務。每個任務都有對應的狀態(Start、Loading、Loaded、Executing 和 Executed)。一個計算設備上的 Loader 和 Executor 共享同一個任務隊列。

| 查詢調度機制

Milvus 查詢任務調度原理

1)Milvus server 啟動時,會根據配置文件中的 gpu_resource_config 啟動對應的 GpuResource,DiskResource 和 CpuResource 啟動暫時不可以通過配置修改。啟動的 Resource 是配置文件中 search_resources 和 build_index_resources 的並集,在這裡是{gpu0, gpu1}。

<code>gpu_resource_config:  
enable: true
cache_capacity: 4
search_resources:
- gpu0
- gpu1
build_index_resources:   
- gpu0/<code>


Milvus 查詢任務調度原理

2)收到請求。前面文章提到過,關於 Table 的一些元數據(Metadata)存在一個外部的數據庫中,單機場景下是 SQLite 或者 MySQL,分佈式場景下是 MySQL。所以 Milvus server 在收到一條搜索請求後會驗證 Table 是否存在、Dimension 是否一致,然後讀出此 Table 的 TableFile 列表。


Milvus 查詢任務調度原理

3)創建 SearchTask。因為每個的 TableFile 的計算是獨立進行,與其他 TableFile 無關的。所以,Milvus 為每一個 TableFile 創建一個 SearchTask。SearchTask 上包含搜索的目標向量、搜索參數以及 TableFile 的文件名。SearchTask 是任務調度的基本單位。

Milvus 查詢任務調度原理

4)選擇計算設備。一個 SearchTask 在什麼設備上計算取決於在每個設備上的 “預計完成時間”。”預計完成時間” 指這個任務從現在到計算完成的預計時間。

例如,一個 SearchTask 的數據塊已經加載到內存裡了,而此時 CPU 的任務隊列中存在一個 SearchTask 等待計算,GPU 的計算隊列是空閒狀態,那麼 CPU 的“預計完成時間”等於前一個 SearchTask 的預計消耗時間加上當前 SearchTask 的預計消耗時間,而 GPU 的“預計完成時間”等於數據塊加載到 GPU 的時間加上當前 SearchTask 的預計消耗時間。某個 Resource 上單個 SearchTask 的預計消耗時間等於這個 Resource 上每個 SearchTask 的平均執行時間。然後選擇一個“預計完成時間”最短的設備,分配 SearchTask 在這個設備上計算。

這裡假設 GPU1 的估計計算完成時間比較短。

Milvus 查詢任務調度原理

5)將 SearchTask 加入到 DiskResource 的任務隊列上。

6)SearchTask 被移動到 CpuResource 的任務隊列上。CpuResource 上的加載線程會依次完成任務隊列中每一個任務的數據加載。CpuResource 把對應的數據塊文件讀入內存。

7)SearchTask 被移動到 GpuResource 上。GpuResource 上的加載線程同樣會依次完成數據從內存到顯存的拷貝。GpuResource 把對應的數據塊從內存讀入顯存。

8)SearchTask 在 GpuResource 上被執行。由於單個 SearchTask 的計算結果一般不會特別大,所以在這一步直接就將結果傳輸回到了內存。

Milvus 查詢任務調度原理

9)SearchTask 的結果與整個搜索請求的結果進行歸併。


Milvus 查詢任務調度原理

10)所有 SearchTask 都完成後,Milvus 會把整個搜索的結果返回給客戶端。

索引構建機制

索引構建流程跟查詢流程基本一致,只不過沒有最後的歸併階段,就不在此贅述了。


| 性能優化

緩存

前面提到,數據塊在計算之前需要被加載到對應的存儲設備上,如內存、顯存。為了避免重複的數據加載,Milvus 引入了基於 LRU(Least Recently Used,最近最少使用)的緩存。當緩存充滿後,新的數據塊會把老的數據塊“擠”出去。用戶可以根據當前系統存儲空間的大小,通過配置文件設置緩存的大小。我們推薦設置足夠大的緩存空間存儲待查向量數據,從而有效減少數據加載時間,提高系統搜索性能。

數據加載與計算重疊

緩存並不能滿足我們所以對於搜索性能的追求,內存不夠大、數據規模大等因素都可能導致數據需要重新加載。我們需要儘可能減少加載數據時間對搜索性能的影響。由於數據加載無論是硬盤加載到內存還是內存加載到顯存,都是 IO 操作,不怎麼佔用處理器的計算資源,所以我們考慮並行地做數據的加載和計算以提高資源利用率。

我們將一個數據塊的計算分成了3個階段(磁盤加載到內存,CPU 計算,結果歸併)或4個階段(磁盤加載到內存,內存加載到顯存,GPU 計算且結果拷回,結果歸併)。以3階段為例,我們可以啟動3個線程分別負責做這3件事情以達到類似 Instruction pipelining 的效果。由於結果集大多數情況下比較小,結果歸併佔用時間並不多,在一些情況下,計算與數據加載重疊能使得整個查詢時間降到原查詢時間的約1/2。

Milvus 查詢任務調度原理


| 問題與解決方案

不一樣的傳輸速度

Milvus 早先的多卡任務分發策略是以 Round Robin 的方式分發。這個分發策略在我們一臺4卡的服務器上工作地很好,查詢性能比單卡提升約4倍,但是在我們的雙卡的開發機上卻沒有2倍的性能提升。關於這個問題我們還做了一些實驗,發現對於某一張卡的數據拷貝速度是11GB/s,而另一張卡只有3GB/s。在查閱了主板相關資料後最終確認,主板上和兩張顯卡的連接分別是 PCIe x16 和 PCIe x4,也就是說這兩張卡的拷貝速度就會不一樣。於是後來我們加入了”拷貝時間”來衡量每個 SearchTask 的最優計算設備。


| 未來的工作

更復雜的硬件環境

實際情況下的硬件環境可以更加複雜,像多路 CPU、NUMA 架構的內存、NVLink 和NVSwitch 這類的卡間通信等這些硬件環境給系統留下了很多的優化空間。

查詢的優化

在實驗中我們發現了一些性能點,例如在 Server 同時接收到很多個對於同一張表的查詢時,這些查詢在一定情況下是可以合併的,利用好數據局部性可以使系統獲得更好的性能。此類的優化我們會在未來的開發過程中慢慢實現。




到這裡,我們已經瞭解了查詢任務在單機多卡的場景下是如何被調度和執行的。Milvus 內部更多的實現原理我們會在其他文章中給出。


分享到:


相關文章: