作者:陳裘凱( 求索)
前言
KubeDL 是阿裡開源的基于 Kubernetes 的 AI 工作負載管理架構,取自"Kubernetes-Deep-Learning"的縮寫,希望能夠依托阿裡巴巴的場景,将大規模機器學習作業排程與管理的經驗反哺社群。目前 KubeDL 已經進入 CNCF Sandbox 項目孵化,我們會不斷探索雲原生 AI 場景中的最佳實踐,助力算法科學家們簡單高效地實作創新落地。
在最新的 KubeDL Release 0.4.0 版本中,我們帶來了模型版本管理(ModelVersion)的能力,AI 科學家們可以像管理鏡像一樣輕松地對模型版本進行追蹤,打标及存儲。更重要的是,在經典的機器學習流水線中,“訓練”與“推理”兩個階段相對獨立,算法科學家視角中的“訓練->模型->推理”流水線缺乏斷層,而“模型”作為兩者的中間産物正好能夠充當那個“承前啟後”的角色。
Github:
https://github.com/kubedl-io/kubedl網站:
https://kubedl.io/model/intro/模型管理現狀
模型檔案是分布式訓練的産物,是經過充分疊代與搜尋後保留的算法精華,在工業界算法模型已經成為了寶貴的數字資産。通常不同的分布式架構會輸出不同格式的模型檔案,如 Tensorflow 訓練作業通常輸出 CheckPoint(*.ckpt)、GraphDef(*.pb)、SavedModel 等格式,而 PyTorch 則通常以 .pth 字尾,不同的架構會在加載模型時解析其中承載的運作時的資料流圖、運作參數及其權重等資訊,對于檔案系統來說,它們都是一個(或一組)特殊格式的檔案,就像 JPEG 和 PNG 格式的圖像檔案一樣。
是以典型的管理方式就是把它們當作檔案,托管在統一的對象存儲中(如阿裡雲 OSS 和AWS S3),每個租戶/團隊配置設定一個目錄,各自成員再把模型檔案存儲在自己對應的子目錄中,由 SRE 來統一進行讀寫權限的管控:

這種管理方式的優缺點都很明顯:
- 好處是保留了使用者的 API 使用習慣,在訓練代碼中将自己的目錄指定為輸出路徑,之後将雲存儲的對應目錄 mount 到推理服務的容器内加載模型即可;
- 但這對 SRE 提出了較高的要求,不合理的讀寫權限授權及誤操作,可能造成檔案權限洩露,甚至大面積的誤删除;同時基于檔案的管理方式不易實作模型的版本管理,通常要求使用者自身根據檔案名來标記,或是上層平台自己承擔版本管理的複雜度;此外,模型檔案與算法代碼/訓練參數的對應關系也無法直接映射,甚至同個檔案在多次訓練中會被多次覆寫,難以追溯曆史;
基于以上現狀,KubeDL 充分結合了 Docker 鏡像管理的優勢,引入了一套 Image-Based 的鏡像管理 API,讓分布式訓練和推理服務結合得更緊密自然,同時也極大簡化了模型管理的複雜度。
從鏡像出發
鏡像(Image)是 Docker 的靈魂,也是容器時代的核心基礎設施。鏡像本身即分層的不可變檔案系統,模型檔案天然地可以作為其中的一個獨立鏡像層,兩者結合的還會迸發出其他火花:
- 使用者不用再面向檔案管理模型,而是直接使用 KubeDL 提供的 ModelVersion API 即可,訓練與推理服務之間通過 ModelVersion API 橋接;
- 與鏡像一樣,可以對模型打 Tag 實作版本追溯,并推送到統一的鏡像 Registry 存儲,通過 Registry 進行鑒權,同時鏡像 Registry 的存儲後端還可以替換成使用者自己的 OSS/S3,使用者可以平滑過渡;
- 模型鏡像一旦建構完畢,即成為隻讀的模闆,無法再被覆寫及篡寫,踐行 Serverless “不可變基礎設施” 的理念;
- 鏡像層(Layer)通過壓縮算法及哈希去重,減少模型檔案存儲的成本并加快了分發的效率;
在“模型鏡像化”的基礎上,還可以充分結合開源的鏡像管理元件,最大化鏡像帶來的優勢:
- 大規模的推理服務擴容場景中,可以通過 Dragonfly 來加速鏡像分發效率,面對流量突發型場景時可以快速彈出無狀态的推理服務執行個體,同時避免了挂載雲存儲卷可能出現的大規模執行個體并發讀時的限流問題;
- 日常的推理服務部署,也可以通過 OpenKruise 中的 ImagePullJob 來提前對節點上的模型鏡像進行預熱,提升擴容釋出的效率。
Model 與 ModelVersion
KubeDL 模型管理引入了 2 個資源對象:Model 及 ModelVersion,Model 代表某個具體的模型,ModelVersion 則表示該模型疊代過程中的一個具體版本,一組 ModelVersion 從同一個 Model 派生而來。以下是示例:
apiVersion: model.kubedl.io/v1alpha1
kind: ModelVersion
metadata:
name: my-mv
namespace: default
spec:
# The model name for the model version
modelName: model1
# The entity (user or training job) that creates the model
createdBy: user1
# The image repo to push the generated model
imageRepo: modelhub/resnet
imageTag: v0.1
# The storage will be mounted at /kubedl-model inside the training container.
# Therefore, the training code should export the model at /kubedl-model path.
storage:
# The local storage to store the model
localStorage:
# The local host path to export the model
path: /foo
# The node where the chief worker run to export the model
nodeName: kind-control-plane
# The remote NAS to store the model
nfs:
# The NFS server address
server: ***.cn-beijing.nas.aliyuncs.com
# The path under which the model is stored
path: /foo
# The mounted path inside the container
mountPath: /kubedl/models
---
apiVersion: model.kubedl.io/v1alpha1
kind: Model
metadata:
name: model1
spec:
description: "this is my model"
status:
latestVersion:
imageName: modelhub/resnet:v1c072
modelVersion: mv-3
Model 資源本身隻對應某類模型的描述,并追蹤最新的版本的模型及其鏡像名告知給使用者,使用者主要通過 ModelVersion 來自定義模型的配置:
- modelName:用來指向對應的模型名稱;
- createBy:建立該 ModelVersion 的實體,用來追溯上遊的生産者,通常是一個分布式訓練作業;
- imageRepo:鏡像 Registry 的位址,建構完成模型鏡像後将鏡像推送到該位址;
- storage:模型檔案的存儲載體,目前我們支援了 NAS,AWSEfs 和 LocalStorage 三種存儲媒體,未來會支援更多主流的存儲方式。以上的例子中展示了兩種模型輸出的方式(本地存儲卷和 NAS 存儲卷),一般隻允許指定一種存儲方式。
當 KubeDL 監聽到 ModelVersion 的建立時,便會觸發模型建構的工作流:
- 監聽 ModelVersion 事件,發起一次模型建構;
- 根據 storage 的類型建立出對應的 PV 與 PVC 并等待 volume 就緒;
- 建立出 Model Builder 進行使用者态的鏡像建構,Model Builder 我們采用了 kaniko 的方案,其建構的過程與鏡像格式,和标準的 Docker 完全一緻,隻不過這一切都在使用者态發生,不依賴任何主控端的 Docker Daemon;
- Builder 會從 volume 的對應路徑中拷貝出模型檔案(可以是單檔案也可以是一個目錄),并将其作為獨立的鏡像層來建構出一個完整的 Model Image;
- 把産出的 Model Image 推送到 ModelVersion 對象中指定的鏡像 Registry 倉庫;
- 結束整個建構過程;
至此,該 ModelVersion 對應版本的模型便固化在了鏡像倉庫中,可以分發給後續的推理服務進行消費。
從訓練到模型
雖然 ModelVersion 支援獨立建立并發起建構,但我們更期望在分布式訓練作業成功結束後自動觸發模型的建構,天然串聯成一個流水線。
KubeDL 支援這種送出方式,以 TFJob 作業為例,在發起分布式訓練時即指定好模型檔案的輸出路徑和推送的倉庫位址,當作業成功執行完畢時就會自動建立出一個 ModelVersion 對象,并将 createdBy 指向上遊的作業名,而當作業執行失敗或提前終止時并不會觸發 ModelVersion 的建立。
以下是一個分布式 mnist 訓練的例子,其将模型檔案輸出到本地節點的
/models/model-example-v
1
路徑,當順利運作結束後即觸發模型的建構:
apiVersion: "training.kubedl.io/v1alpha1"
kind: "TFJob"
metadata:
name: "tf-mnist-estimator"
spec:
cleanPodPolicy: None
# modelVersion defines the location where the model is stored.
modelVersion:
modelName: mnist-model-demo
# The dockerhub repo to push the generated image
imageRepo: simoncqk/models
storage:
localStorage:
path: /models/model-example-v1
mountPath: /kubedl-model
nodeName: kind-control-plane
tfReplicaSpecs:
Worker:
replicas: 3
restartPolicy: Never
template:
spec:
containers:
- name: tensorflow
image: kubedl/tf-mnist-estimator-api:v0.1
imagePullPolicy: Always
command:
- "python"
- "/keras_model_to_estimator.py"
- "/tmp/tfkeras_example/" # model checkpoint dir
- "/kubedl-model" # export dir for the saved_model format
% kubectl get tfjob
NAME STATE AGE MAX-LIFETIME MODEL-VERSION
tf-mnist-estimator Succeeded 10min mnist-model-demo-e7d65
% kubectl get modelversion
NAME MODEL IMAGE CREATED-BY FINISH-TIME
mnist-model-demo-e7d65 tf-mnist-model-example simoncqk/models:v19a00 tf-mnist-estimator 2021-09-19T15:20:42Z
% kubectl get po
NAME READY STATUS RESTARTS AGE
image-build-tf-mnist-estimator-v19a00 0/1 Completed 0 9min
通過這種機制,還可以将其他“僅當作業執行成功才會輸出的 Artifacts 檔案”一起固化到鏡像中,并在後續的階段中使用。
從模型到推理
有了前面的基礎,在部署推理服務時直接引用已建構好的 ModelVersion,便能加載對應模型并直接對外提供推理服務。至此,算法模型生命周期(代碼->訓練->模型->部署上線)各階段通過模型相關的 API 聯結了起來。
通過 KubeDL 提供的 Inference 資源對象部署一個推理服務時,隻需在某個 predictor 模闆中填充對應的 ModelVersion 名,Inference Controller 在建立 predictor 時會注入一個 Model Loader,它會拉取承載了模型檔案的鏡像到本地,并通過容器間共享 Volume 的方式把模型檔案挂載到主容器中,實作模型的加載。如上文所述,與 OpenKruise 的 ImagePullJob 相結合我們能很友善地實作模型鏡像預熱,來為模型的加載提速。為了使用者感覺的一緻性,推理服務的模型挂載路徑與分布式訓練作業的模型輸出路徑預設是一緻的。
apiVersion: serving.kubedl.io/v1alpha1
kind: Inference
metadata:
name: hello-inference
spec:
framework: TFServing
predictors:
- name: model-predictor
# model built in previous stage.
modelVersion: mnist-model-demo-abcde
replicas: 3
batching:
batchSize: 32
template:
spec:
containers:
- name: tensorflow
args:
- --port=9000
- --rest_api_port=8500
- --model_name=mnist
- --model_base_path=/kubedl-model/
command:
- /usr/bin/tensorflow_model_server
image: tensorflow/serving:1.11.1
imagePullPolicy: IfNotPresent
ports:
- containerPort: 9000
- containerPort: 8500
resources:
limits:
cpu: 2048m
memory: 2Gi
requests:
cpu: 1024m
memory: 1Gi
對于一個完整的推理服務,可能同時 Serve 多個不同模型版本的 predictor,比如在常見搜尋推薦的場景中,期望以 A/B Testing 實驗來同時對比多次模型疊代的效果,通過 Inference+ModelVersion 可以很容易做到。我們對不同的 predictor 引用不同版本的模型,并配置設定合理權重的流量,即可達到一個推理服務下同時 Serve 不同版本的模型并灰階比較效果的目的:
apiVersion: serving.kubedl.io/v1alpha1
kind: Inference
metadata:
name: hello-inference-multi-versions
spec:
framework: TFServing
predictors:
- name: model-a-predictor-1
modelVersion: model-a-version1
replicas: 3
trafficWeight: 30 # 30% traffic will be routed to this predictor.
batching:
batchSize: 32
template:
spec:
containers:
- name: tensorflow
// ...
- name: model-a-predictor-2
modelVersion: model-version2
replicas: 3
trafficWeight: 50 # 50% traffic will be roted to this predictor.
batching:
batchSize: 32
template:
spec:
containers:
- name: tensorflow
// ...
- name: model-a-predictor-3
modelVersion: model-version3
replicas: 3
trafficWeight: 20 # 20% traffic will be roted to this predictor.
batching:
batchSize: 32
template:
spec:
containers:
- name: tensorflow
// ...
總結
KubeDL 通過引入 Model 和 ModelVersion 兩種資源對象,與标準的容器鏡像相結合實作了模型建構,打标與版本追溯,不可變存儲與分發等功能,解放了粗放型的模型檔案管理模式,鏡像化還可以與其他優秀的開源社群相結合,實作鏡像分發加速,模型鏡像預熱等功能,提升模型部署的效率。同時,模型管理 API 的引入很好地連接配接了分布式訓練與推理服務兩個原本割裂的階段,顯著提升了機器學習流水線的自動化程度,以及算法科學家上線模型、實驗對比的體驗和效率。我們歡迎更多的使用者試用 KubeDL,并向我們提出寶貴的意見,也期待有更多的開發者關注以及參與 KubeDL 社群的建設!
KubeDL Github 位址:
戳
此處,立即了解 KubeDL 項目!