TensorFlow Serving的原理和程式碼實現
本文介紹tensorflow Serving的原理和程式碼實現, 並提供簡要的程式碼閱讀指導.
如何serve一個模型
具體的步驟可以參考官方文件. 主要包括兩個部分：
1. 匯出模型
1. 啟動服務
需要說明的是匯出模型部分. 如果要把我們訓練的模型拿來提供服務, 除了模型本身外, 還需要一些額外的資訊, 比如模型的名稱, 輸入、輸出對應的tensor資訊, 方法名, 這些東西可以讓TFS進行請求資料的格式檢查以及目標模型查詢. 這就是模型匯出的作用. 直接拿一個checkpoint檔案之類的是不能用的. TF使用SavedModel格式匯出模型, 並提供了相關的工具(tf.saved_model.builder.SavedModelBuilder).
TFS的功能
- 支援多種模型服務策略,比如用最新版本/所有版本/指定版本, 以及動態策略更新、模型的增刪等
- 自動載入/解除安裝模型
- Batching
- 多種平臺支援(非TF平臺)
TFS的架構
本節簡要介紹各模組的主要功能, 後續章節介紹他們相互之間是如何協作的.
Servable
Servable是對模型的抽象, 但是任何能夠提供演算法或者資料查詢服務的實體都可以是Servable, 不一定是機器學習模型.
在我們常用的場景下, Servable就是模型. 所以本文有時會混用模型和Servable.
ServerCore
整個服務系統的建立維護, 建立http rest server、grpc server和模型管理部分(AspiredVersionManager)之間的關係等.
AspiredVersionManager
模型管理的上層控制部分. 負責執行Source發出的模型管理指令. 一部分功能通過回撥的方式由Source呼叫, 一部分由獨立執行緒執行.
BasicManager
負責Servable的管理, 包括載入、解除安裝、狀態查詢、資源跟蹤等, 對外提供如下介面操作Servable
- ManageServable
- LoadServable
- UnloadServable
- StopManagerServable
另外提供介面查詢servableHandle(GetUntypedServableHandle), 也就是載入好的模型,供http rest或者grpc server呼叫進行推理.
所有受管理的servable都放在ManagedMap裡面, 已經正常載入的servable同時也放在ServingMap裡進行管理, 提供查詢介面.
LoaderHarness
LoaderHarness對Loader提供狀態跟蹤, ServingMap和ManagedMap裡面儲存的都是LoaderHarness物件,只有通過LoaderHarness才能訪問Loader的介面.
Loader
Loader對Servable的生命週期進行控制, 包括load/unload介面,資源預估介面等. 載入之後的Servable也存在Loader裡面.
Adapter
Adapter是為了將Source(比如檔案系統)轉換成Loader而引入的抽象, 這樣server core的實現和具體的平臺解耦.
SourceRouter
Adapter是平臺相關的, 每個平臺一個Adapter, 但是Source是和Servable相關的, 這樣在Adapter和Source之間存在一對多的關係, Router負責維護這些對應關係.
Source
Source是對Servable的來源(Source)的抽象, 比如模型檔案是某個模型的Source. Source監控外部的資源（如檔案系統）, 發現新的模型版本, 並通知Target.
Target
Target是和Source對應的抽象概念, AspiredVersionManager、Router都是Target.
啟動過程
TFS啟動的全部引數可以參考main.c. 主要的引數包括服務埠（gprc和http rest埠）和模型配置. 其中模型配置可以直接指定（名稱（model_name）、路徑（model_base_path）等）, 也可以使用檔案指定（model_config_file,格式參考model_server_config.proto）. 如果只啟動單個模型的服務可以使用引數指定, 如果是多個模型必須使用檔案. 其他的引數可以使用預設值.
啟動過程主要是建立ServerCore物件, 並啟動grpc server和http server.
ServerCore物件可以認為是系統中樞, 模型的維護, 服務請求的處理都是由他完成. ServerCore通過BasicManager管理所有的model(多版本號), 並查處模型已經提供預測、分類、迴歸請求.
ServerCore啟動的時候建立AspiredVersionManager, AspiredVersionManager會啟動定時任務(執行緒), 用於處理AspiredVersionRequest訊息, 其實就是模型的載入、解除安裝.
啟動的時候ServerCore還會根據模型配置建立檔案系統掃描任務, 定時掃描模型檔案目錄並進行相應的處理
http rest服務啟動後, 會監聽http post請求, 將請求（json）轉換成protobuf格式的訊息, 通過serverCore查詢對應的模型版本, 獲取對應的已載入的模型, 進行運算並返回結果.
rgpc服務與 http rest服務類似.
模型維護
檔案系統掃描
Source是TFS定義的對未載入模型物件的抽象, 目前實現了兩種Source, 一種是StaticStoragePathSource,一種是FileSystemStoragePathSource. 前者是簡單的靜態的模型檔案儲存系統, 僅僅在啟動時觸發模型的載入, 沒有其他動作. 後者是動態的Source, 能監測儲存系統的變化併發出通知.
TFS實現Source時將模組職責劃分的很清晰, Source的職責就是監測變化, 如何處理則由Source的使用者決定, 所以Source有一個介面SetAspiredVersionsCallback, 可以設定回撥函式用於通知AspiredVersion的變化. Source在變化的時候就會呼叫設定的回撥函式.
作為Source的對等物件, 系統也定義了Target, 有介面GetAspiredVersionsCallback, 用於獲取處理AspiredVersions的回撥介面, 然後我們就可以將Target和Source連起來了.
template <typename T>
void ConnectSourceToTarget(Source<T>* source, Target<T>* target) {
source->SetAspiredVersionsCallback(target->GetAspiredVersionsCallback());
}Source和ServerCore的關係是這樣的
Source --> Router --> Adapter --> AspiredVersionManager上述連線關係裡面, Router和Adapter既是Source又是Target, AspiredVersionManager是Target. 但是Router沒有實現Source介面, 而是要求在建立Router物件時直接將Adapter作為引數, 這樣實現主要目的是建立一對多的關係.
系統根據所支援平臺的個數（tensorflow算是一種平臺）建立Adapter, 一種平臺對應一個Adapter, 負責建立模型載入器Loader. 對於tensorflow平臺, 對應的adapter是SavedModelBundleSourceAdapter.
Router負責根據模型名稱查詢對應的平臺(model_config裡面有指定平臺名稱), 從而定位到對應的Adapter.
這些連線關係是在系統啟動, 或者更新model-config的時候建立的.
預設配置下, FileSystemStoragePathSource為Source的例項, SavedModelBundleSourceAdapter為Adapter的例項, DynamicSourceRouter為Router的例項.
- FileSystemStoragePathSource有自己單獨的工作執行緒, 週期查詢檔案系統, 發現每個模型的版本, 根據指定的servable_version_policy(model_config), 建立ServableData(模型名, 版本號, 路徑), 傳給Router
- Router根據路由找到對應的adapter, 傳給Adataper
- Adapter將ServableData(模型名, 版本號, 路徑)轉換成ServableData(模型名, 版本, Loader), 傳給AspiredVersionManager
- AspiredVersionManager將這些資訊存到pending_aspired_versions_requests_, 等待另外一個工作執行緒(AspiredVersionsManager_ManageState_Thread)處理
上述訊息傳遞的方式是依次呼叫下游的SetAspiredVersions函式.
模型載入/解除安裝
上節提到的工作執行緒ManageState_Thread是在AspiredVersionsManager建立的時候啟動的定時執行緒, 負責處理pending_aspired_versions_requests_裡面的ServableData.
manage_state_thread_.reset(new PeriodicFunction(
[this]() {
this->FlushServables();
this->HandlePendingAspiredVersionsRequests();
this->InvokePolicyAndExecuteAction();
},
manage_state_interval_micros));該執行緒的工作分3部分, 如上述程式碼所示
FlushServables主要目的是將異常狀態的模型清理掉, 或者停止載入.
HandlePendingAspiredVersionsRequests取出每個模型的資訊分別處理, 如果發現當前要載入的模型版本已經存在, 需要等待之前的模型完成服務並退出, 這叫re-aspired version. 如果不是這種情況, 計算需要載入的模型和需要解除安裝的模型, 將新載入的模型管理起來（加到管理列表）,將需要解除安裝的模型打上標記並停止其載入.
InvokePolicyAndExecuteAction每次只會執行一個模型的一個動作（load/unload). 具體方法是每個模型根據aspired_version_policy（AvailabilityPreservingPolicy/ResourcePreservingPolicy）選擇一個動作, 然後所有模型的選擇動作放在一起排序, unload優先load, 決定處理哪一個模型. 執行動作的時候, 會呼叫Loader的相應函式, 並設定相關的狀態(參考模型狀態管理).
可以看出ManageState_Thread並不是一股腦的進行模型的載入、解除安裝等操作, 而是兼顧了資源佔用、服務可用性、系統負荷的, 考慮周到.
模型管理
AspiredVersionManager的成員BasicManager負責模型管理, 把一個模型版本加入到BasicManager裡面就叫管理一個模型(manager a model). BasicManager通過LoaderHarness間接管理模型, LoaderHarness管理一個模型的生命週期, 持有模型的Loader物件, 在合適的時候呼叫Loader的Load/Unload完成狀態遷移.
LoaderHarness有自己的狀態記錄, 每次執行動作時都會進行狀態遷移.
進行模型狀態管理的同時, BasicManager還會將模型的服務狀態釋出到EventBus(servable_event_bus_), 便於其他模組對這些狀態變化進行訂閱.
快速模型載入
首次啟動的時候, 採用快速載入模式, 實現方法是臨時增加模型載入執行緒（4倍可用cpu個數）. 完成載入後恢復執行緒數.
程式碼
// The number of load threads used to load the initial set of models at
// server startup. This is set high to load up the initial set of models
// fast, after this the server uses num_load_threads.
int32 num_initial_load_threads = 4.0 * port::NumSchedulableCPUs();Aspired version policy
AspiredVersionPolicy是用來決定一個模型的多個版本誰先處理, AvailabilityPreservingPolicy的目標是保證服務可用, 會臨時犧牲一些資源, 而
ResourcePreservingPolicy是優先保證佔用更少的資源, 可能會犧牲服務可用性.
程式碼裡面的註釋提供了很好的解釋, 可參考.
Servable Version Policy
ServableVersionPolicy
定義了模型的多個版本如何進行選擇. 注意和Aspired Version Policy的關係, 一個是如何選擇版本, 一個是選擇了版本後, 如何選擇執行先後順序.
目前提供3種方式：
- all 所有的版本
- latest 最新的N個版本
- specific 一個或一些指定的版本號
模型配置動態更新
tfs的main預設並沒有提供模型配置檔案的動態更新, 但是呼叫ServerCore::ReloadConfig(const ModelServerConfig& new_config)就可以完成更新. 可以自己包裝該介面在合適的時間進行呼叫.
動態更新可以增加、刪除模型, 修改版本策略（比如從最新版本到指定某版本）等.
動態更新的實現也不復雜, 更新Router的路由策略, 更新Source就可以了. 程式碼實現可以參考ServerCore::AddModelsViaModelConfigList
TFS目前還不支援動態增加平臺.
資源管理
TFS目前僅支援記憶體資源的管理，類ResourceTracker用來跟蹤當前Servable消耗的總資源，當有新的Servable需要載入的時候，會計算剩下的資源是否夠用, 並預留資源(BasicManager::ReserveResources).
ServerCore.Options.total_model_memory_limit_bytes控制總資源，預設設定無上限.
SavedModelBundleFactory提供了對TF模型資源的評估方法，簡單的將模型檔案大小乘1.2倍(程式碼, EstimateResourceFromPath).
Batching
Batching是提高服務效能的一個有效辦法, 最簡單的batching就是把多個單獨請求打包到一起, 由TF Session一次運算得出結果.
Batching的實現就是在普通的Tensorflow Session之外包裝一個BatchingSession，負責快取、排程. 參考程式碼SavedModelBundleFactory::Create. Batching分兩個部分實現，BatchScheduler和BatchingSession.
BatchingSession對外提供Run和ListDevices介面. 有Run請求的時候, 將請求打包成一個BatchTask，交給BatchScheduler去處理，並等待處理結束，取出結果返回.
BatchScheduler是一個底層的排程器，擁有自己的執行緒池, 負責將多個BatchTask合併處理.
Model Warmup
模型載入後，如果需要Warmup，從模型檔案目錄中取出預先存好的請求資料，呼叫模型進行推理，如此可以將模型”熱身”，避免首次處理服務請求時時延過大。
配置
平臺配置
如果需要使用分散式TF伺服器, 可以在這裡指定. 預設情況下使用本地Session(DirectSession)
- Session 引數
- Batching引數
模型配置
與Tensorflow的關係
TFS對TF是原始碼級別的依賴, 兩者的版本號保持一致, TFS在載入模型、執行推理的過程中, 都是呼叫TF的庫. TFS使用的很多基本構件, 比如多執行緒庫/BatchScheduler, 都是直接使用TF的程式碼.
效能
TFS在如下方面做出了效能提升的設計：
- Batching
- Fast Model Loading
- Model Warmup
- Availability/Resource Proserving Policy