【第四期】AI Talk：AI 計算，從單機到集羣（中）
- 引言
- Kubernetes
- Kubernetes 的 GPU 資源管理與調度
- Device Plugin
- 準備工作
- 設備註冊
- 設備分配
- 設備回收
- Device Manager
- 用戶使用
- 回顧與展望
- 總結
- 參考
引言
在 AI 計算：從單機到集羣（上）中我們介紹了任務計算容器化的概念，解決單機計算的便利性問題。當我們進一步深入使用，尤其是有多臺計算節點時，會面臨一個新問題：多臺 GPU 節點如何實現 GPU 資源的合理自動分配，達到共享使用 GPU 集羣資源的目的。本篇介紹如何利用 Kubernetes 實現 GPU 資源的集羣化管理，我們結合 Kubernetes 相關源碼，詳細分析 Kubernetes 框架下 GPU 設備管理邏輯和方式，以及如何實現對 GPU 資源的調度，幫助大家邁出從單機到集羣的關鍵一步。
Kubernetes
容器技術簡化了計算框架單機運行的問題，但是沒法解決多機資源分配問題。雖然目前很多計算框架支持分散式運行，但缺少資源分配和調度的功能。這就是本篇我們所需要解決的問題，我們期望計算平臺（集羣）能夠提供 GPU 任務調度、監控、失敗重啟等全生命週期管理的功能。當集羣計算規模擴大時，如果沒有這些功能，我們很難手工地去每一個計算節點上啟動計算任務，也無法實時監控任務運行。除此之外，目前大多數分散式計算框架不支持生命週期管理，結束的訓練進程並不會自動關閉，這也需要進行額外的處理。同時，當多用戶共享使用計算資源時，如果依靠人工協調資源分配，會帶來集羣資源利用效率低下和使用繁瑣等問題。總之，當存在多機集羣或者多用戶共享使用的情況時，我們需要一種平臺幫助我們實現以下目標：
- GPU 計算資源的自動調度與管理
- 計算任務的全生命週期管理
- 自動實現多用戶任務共享使用計算資源
而我們本篇將要介紹的 Kubernetes 能夠滿足我們上述要求。Kubernetes (簡稱：k8s) 是 Google 開源的容器集羣管理系統（內部代號：Borg）。Kubernetes 是一套完備的分散式系統平臺，具有完備的集羣管理能力，以容器技術為基礎，為容器化的應用提供部署運行、資源調度、服務發現和動態伸縮等一系列完整功能，提高了大規模容器集羣管理的便捷性，包括：容器自動化部署、可擴展資源自動調度機制以及多粒度的資源配額管理等。這是 Kubernetes 針對傳統雲服務和雲計算所提供的一攬子功能。然而要實現我們所需要的目標，僅僅具備這些功能是不夠的。我們需要能夠對 GPU 資源進行管理和調度，包括更進一步的配額管理功能等。
Kubernetes 從 1.6 版本開始，逐漸開始支持 GPU 資源的調度，而且功能實現也在快速更新迭代，大致分為兩個階段：
- 在 1.8 版本之前，GPU 的管理在功能實現上比較簡單粗暴，處於試驗階段，具備初步的基本功能，但還不能大規模的應用在生產環境。
- 在 1.8 版本之後，Kubernetes 重構了 GPU 的管理邏輯，引入 Device Manager 和 Device Plugin 組件，以更加規範完善的框架實現對擴展硬體的支持，不侷限於只支持 GPU 硬體資源，還支持 FPGA，InfiniteBand 等。除此之外，還添加了設備健康檢查等新功能。
本篇以 Kubernetes 1.10 版本為例，介紹 Kubernetes 如何實現 GPU 資源管理與調度，以及用戶如何利用 Kubernetes 提交 GPU 任務。
Kubernetes 的 GPU 資源管理與調度
Kubernetes 1.8 版本之後，由 Device Manager 管理擴展硬體資源，包括 GPU，FPGA，InfiniteBand 等。不同硬體資源或者同一種硬體的不同廠商均可開發對應設備的 Device Plugin。Device Plugin 按照約定的 API 與 Kubernetes 中的 Device Manager 進行交互，實現對設備的管理、分配、回收以及運行狀態監控和健康檢查。
我們通過 Kubernetes 的代碼可以發現，Device Manager 的功能邏輯集成在 Kubernetes 的核心組件 kubelet 代碼中。Device Manager 與 Device Plugin 交互方式如下圖所示。圖中左側是 Device Manager 的功能邏輯組件（綠色），右側是 Device Plugin 的功能邏輯組件（紅色），其中 Device Plugin 的紅色色塊代表功能執行，Device Manager 的綠色色塊代表功能發起，空心矩形框為雙方的 gRPC Server 邏輯功能。雙方的 gPRC 交互過程分為四種，具體包括：
- 設備註冊
- 設備監控
- 設備分配
- 設備回收
我們以 Nvidia GPU 設備為例，詳細介紹 Device Manager 和 Device Plugin 如何通過這四個交互過程實現對 GPU 的資源管理和調度。Nvidia 官方發布了 k8s-device-plugin，我們以此作為 Device Plugin 的示例，下文簡稱插件。
Device Plugin
準備工作
k8s-device-plugin 在運行之前，需要準備好 GPU 驅動，確保 GPU 能夠正常工作。k8s-device-plugin 並不提供 GPU 驅動安裝工作，需要管理員在每一個計算節點提前安裝好 GPU 驅動。
k8s-device-plugin 以容器方式或者直接運行在計算節點上，推薦以 daemon set 方式，插件啟動的 YAML 參考官方推薦配置文件。
k8s-device-plugin 需要和 Nvidia runtime 配合使用，在集羣使用之前，每一個 GPU 計算節點上安裝 Nvidia runtime，並修改 /etc/docker/daemon.json，設置 Nvidia runtime 為默認的 docker runtime，如下所示：
{ "default-runtime": "nvidia", "runtimes": { "nvidia": { "path": "/usr/bin/nvidia-container-runtime", "runtimeArgs": [] } }}
Kubernetes 版本 1.8、1.9 需要添加 kubelet 參數 --feature-gates=DevicePlugins=true 開啟 Device Plugin 功能，在 1.10 版本，此功能已經默認開啟，無需再添加 kubelet 參數。
設備註冊
Device Manager 內部維護一個供插件註冊的 gRPC server，k8s-device-plugin 利用 gRPC 協議，並通過 /var/lib/kubelet/device-plugins/kubelet.sock 向 Device Manager 發起註冊請求，對應的是圖中紅色色塊的 Register 向左邊的 Registry gRPC Server 發送註冊信息，k8s-device-plugin 註冊過程如下：
- k8s-device-plugin 向 Device Manager 發起 RegisterRequest gRPC 請求，彙報 GPU 設備的信息：
- 設備名稱 http://nvidia.com/gpu
- API 版本號
- k8s-device-plugin 的 gPRC server unix socket: /var/lib/kubelet/device-plugins/nvidia.sock
- Device Manager 響應 RegisterRequest，利用 RegisterResponse 返迴響應結果，如果設備註冊成功後，更新節點擴展設備資源狀態信息。
- 如果 Device Manager 響應成功，k8s-device-plugin 啟動插件端的 gPRC Server，供 Device Manager 與 k8s-device-plugin 通信。
client := pluginapi.NewRegistrationClient(conn)reqt := &pluginapi.RegisterRequest{Version: pluginapi.Version,Endpoint: path.Base(m.socket),ResourceName: resourceName,}_, err = client.Register(context.Background(), reqt)
k8s-device-plugin 通過 kubelet.sock 持續監控 kubelet 的運行狀態，當 kubelet 重啟後，k8s-device-plugin 需要向 Device Manager 重新註冊。由於雙方需要通過 socket 通訊，當 k8s-device-plugin 以容器的方式運行時，需要掛載主機的 /var/lib/kubelet/device-plugins/ 目錄。
設備發現
當插件註冊成功後，Device Manager 通過 ListAndWatch gRPC 請求獲取當前設備的列表和健康狀態，這個交互是雙向的，如果設備狀態發生改變，比如當 Device Plugin 檢測到某個設備不健康的時候，就會主動通知 Device Manager。如果這個不健康的設備處於空閑狀態，Device Manager 就會將其挪出可分配列表。如果該設備已經被某個任務使用，kubelet 中止此任務的使用。
設備發現的功能依賴每個插件根據不同設備的具體情況做出不同的處理。k8s-device-plugin 的設備發現過程調用如下：getDevices 函數獲取 GPU UUID ，並傳遞給 Device Manager。GPU UUID 由 k8s-device-plugin 調用 Nvidia 的 NVML 庫獲取。NVIDIA Management Library (NVML) 是 Nvidia 官方發布的基於 C 語言介面，用於監控和管理 Nvidia GPU 的工具庫。其編譯版隨 GPU 驅動一起發布。Nvidia 的 nvidia-smi 和其他常用的第三方工具均使用 NVML 作為管理 GPU 的底層介面，k8s-device-plugin 由於是基於 go 語言的實現，所以直接使用了 nvidia-docker 1.0 中 NVML go 語言 Binding https://github.com/NVIDIA/nvidia-docker/tree/1.0/src/nvml。
func getDevices() []*pluginapi.Device { n, err := nvml.GetDeviceCount() check(err) var devs []*pluginapi.Device for i := uint(0); i < n; i++ { d, err := nvml.NewDeviceLite(i) check(err) devs = append(devs, &pluginapi.Device{ ID: d.UUID, Health: pluginapi.Healthy, }) } return devs}
設備分配
設備分配由插件實現 Allocate 功能，負責配置硬體環境。kubelet 創建任務時，通過 gPRC 調用插件的 Allocate，完成設備的分配，以及確保設備在容器中能正常使用。常見的操作包括設置環境變數、掛載 volume、初始化容器所需的設備等。
k8s-device-plugin 是如何實現 GPU 設備分配呢？藉助於 AI 計算：從單機到集羣（上）介紹的 Nvidia runtime，極大簡化了 GPU 的分配邏輯，所以 k8s-device-plugin 的 Allocate 實現非常優雅，如下面代碼所示，只需要設置容器的環境變數即可。其具體過程是：
- Device Manager 傳遞待分配 GPU UUID 列表，調用 k8s-device-plugin 的 Allocate 模塊。
- k8s-device-plugin 的 Allocate 模塊設置 Nvidia runtime 的環境變數 NVIDIA_VISIBLE_DEVICES
response := pluginapi.ContainerAllocateResponse{ Envs: map[string]string{ "NVIDIA_VISIBLE_DEVICES": strings.Join(req.DevicesIDs, ","), }, }
設備回收
在設備回收階段，插件可以做一些比如驅動卸載等收尾工作。由於 GPU 不需要每次任務結束時卸載驅動，k8s-device-plugin 無需處理設備回收工作，任務資源的釋放由 kubelet 控制。
Device Manager
Device Manager 的主要功能點包括：
- 提供註冊 gRPC Server， 接受 Device Plugin 的註冊請求，獲取並維護節點上的設備列表。
- 與 Device Plugin 保持長連接通信，為 Device Plugin 提供回調函數，當節點設備狀態改變時，Device Plugin 通知 Device Manager 根據設備列表信息，更新節點設備狀態。
- 設備的分配與管理，維護當前運行任務與已使用設備的映射關係，提供待分配設備列表，並通過調用 Device Plugin 的 Allocate，協助 kubelet 完成任務設備分配相關的工作。
Device Manager 實現結構體如下，包括：
- 通信 socket 文件
- 負責註冊的 gRPC Server
- activePods 獲取當前運行任務
- sourceReady 用於從 checkpoint 移除不活動任務
- callback 提供回調給 Device Plugin，用於設備監控狀態
- healthyDevices 健康設備列表，unhealthyDevices 問題設備列表，allocatedDevices 已使用設備列表
- podDevices 記錄運行任務和已使用的設備映射關係
// ManagerImpl is the structure in charge of managing Device Plugins.type ManagerImpl struct { socketname string socketdir string endpoints map[string]endpoint // Key is ResourceName mutex sync.Mutex server *grpc.Server // activePods is a method for listing active pods on the node // so the amount of pluginResources requested by existing pods // could be counted when updating allocated devices activePods ActivePodsFunc // sourcesReady provides the readiness of kubelet configuration sources such as apiserver update readiness. // We use it to determine when we can purge inactive pods from checkpointed state. sourcesReady config.SourcesReady // callback is used for updating devices states in one time call. // e.g. a new device is advertised, two old devices are deleted and a running device fails. callback monitorCallback // healthyDevices contains all of the registered healthy resourceNames and their exported device IDs. healthyDevices map[string]sets.String // unhealthyDevices contains all of the unhealthy devices and their exported device IDs. unhealthyDevices map[string]sets.String // allocatedDevices contains allocated deviceIds, keyed by resourceName. allocatedDevices map[string]sets.String // podDevices contains pod to allocated device mapping. podDevices podDevices store utilstore.Store pluginOpts map[string]*pluginapi.DevicePluginOptions}
由於篇幅限制，這裡對 Device Manager 的實現細節不多做介紹，展開說一下 Device Manager 如何維護當前運行任務與已使用設備的映射關係。每個計算節點上的 Device Manager 創建保存當前運行任務與已使用設備映射關係的 checkpoint 文件 /var/lib/kubelet/device-plugins/kubelet_internal_checkpoint ，其內容如下，分為兩部分：
- PodDeviceEntries 保存節點正在運行任務的 PodID，ContainerName，ResourceName, 正在使用的 DeviceIDs 列表，以及 Device Plugin 返回的消息 AllocResp。其中，DeviceIDs 中多個 GPU UUID 代表多卡任務，
- RegisteredDevices 保存節點處於健康狀態的設備列表，以 ResourceName 為 key，其值為 DeviceIDs 列表。
將設備使用映射關係通過 checkpoint 文件的方式保存到硬碟上，其目的是為瞭解決由於 kubelet 重啟帶來的設備映射關係信息丟失的問題。當 kubelet 重啟時，自動讀取硬碟上的 checkpoint 文件以獲得重啟前的設備使用映射關係，保證設備映射關係與實際任務使用的一致性，避免將已經分配的設備當做未使用設備重新分配使用的問題。
{ "PodDeviceEntries": [ { "PodUID": "51a38fdb-3ef0-11e8-b8fe-0cc47ae55a2c", "ContainerName": "task1", "ResourceName": "nvidia.com/gpu", "DeviceIDs": [ "GPU-77ccee89-7bbc-8838-a56e-f0ca79518232" ], "AllocResp": "CkIKFk5WSURJQV9WSVNJQkxFX0RFVklDRVMSKEdQVS03N2NjZWU4OS03YmJjLTg4MzgtYTU2ZS1mMGNhNzk1MTgyMzI=" }, { "PodUID": "e7f290e2-3be0-11e8-b8fe-0cc47ae55a2c", "ContainerName": "task2", "ResourceName": "nvidia.com/gpu", "DeviceIDs": [ "GPU-93d815e1-0bda-ea1f-08d9-0864e895553d", "GPU-ffd50dfd-2578-342e-9a53-19b0f3d40852", "GPU-68aed792-9550-6ef7-bd91-f8422efd7b5a", "GPU-a6ec8254-2bd0-3237-142a-496fa2059d73" ], "AllocResp": "Cr4BChZOVklESUFfVklTSUJMRV9ERVZJQ0VTEqMBR1BVLTkzZDgxNWUxLTBiZGEtZWExZi0wOGQ5LTA4NjRlODk1NTUzZCxHUFUtZmZkNTBkZmQtMjU3OC0zNDJlLTlhNTMtMTliMGYzZDQwODUyLEdQVS02OGFlZDc5Mi05NTUwLTZlZjctYmQ5MS1mODQyMmVmZDdiNWEsR1BVLWE2ZWM4MjU0LTJiZDAtMzIzNy0xNDJhLTQ5NmZhMjA1OWQ3Mw==" } ], "RegisteredDevices": { "nvidia.com/gpu": [ "GPU-68aed792-9550-6ef7-bd91-f8422efd7b5a", "GPU-02a18b6f-3098-7c10-33f9-ededd1b150b8", "GPU-e4ce184a-d4a9-8b90-9ba1-9f40ec4cc2d7", "GPU-68258d71-d70e-f8bd-9c9a-b7b5240c8b58", "GPU-a6ec8254-2bd0-3237-142a-496fa2059d73", "GPU-6d8322d9-b8b5-89ce-2804-5cbaacc7b6ef", "GPU-93d815e1-0bda-ea1f-08d9-0864e895553d", "GPU-ffd50dfd-2578-342e-9a53-19b0f3d40852", "GPU-77ccee89-7bbc-8838-a56e-f0ca79518232", "GPU-69bf1feb-66bc-67b9-c38e-c5a38ac93e20" ] }}
用戶使用
我們從用戶角度看一下，從任務提交開始的整個交互工作流程：
- 用戶提交任務申請，通過在 YAML 文件裏指定 http://nvidia.com/gpu 申請 X 個 GPU 卡
- Scheduler 過濾滿足條件的候選節點
- 任務 Pod 被分發到節點，該節點 Device Manager 決定待分配設備的 GPU UUID 列表
- Device Manager 調用 gPRC Allocate，通知 Device Plugin 將 GPU UUID 列表中的設備映射到任務 Pod 中使用（比如上面介紹的：k8s-device-plugin 設置 Nvidia runtime 的環境變數 NVIDIA_VISIBLE_DEVICES）
- 任務完成創建
我們也給出 YAML 文件示例，如下面的 YAML 文件所示，用戶提交 GPU 任務時，只需要通過 http://nvidia.com/gpu 指定 GPU 使用數量，Kubernetes 的 scheduler 查找合適的節點資源，並自動調度到滿足要求的節點上，由節點上的 kubelet 完成任務的啟動和運行。如果沒有資源空餘，則任務會處於 Pending 狀態，等待任務需求的資源滿足，則自動轉入運行狀態。
apiVersion: v1kind: Podmetadata: name: tensorflowspec: containers: - name: tensorflow args: ["sleep 1d"] command: ["/bin/sh", "-c"] image: tensorflow/tensorflow:latest-gpu resources: limits: nvidia.com/gpu: "1"
Kubernetes 支持對不同的 GPU 資源需求做篩選，比如，可以根據 GPU 型號做任務選擇調度。如下所示，指定使用 Tesla P100 卡，則 Kubernetes 會自動調度任務到 P100 卡的節點上。
apiVersion: v1kind: Podmetadata: name: tensorflowspec: containers: - name: tensorflow args: ["sleep 1d"] command: ["/bin/sh", "-c"] image: tensorflow/tensorflow:latest-gpu resources: limits: nvidia.com/gpu: "1" nodeSelector: accelerator: nvidia-tesla-p100
這裡需要進一步說明下使用 k8s-device-plugin 的一個小 bug，由於 GPU 計算節點上的 docker runtime 默認設置為 Nvidia runtime，而 Nvidia runtime 的 NVIDIA_VISIBLE_DEVICES 環境變數默認值為 all。所以，當用戶提交非 GPU 任務時，如下所示，在 YAML 文件中沒有指定 http://nvidia.com/gpu，則此時，在任務容器中能使用該節點上的所有 GPU 資源。這顯然不是我們期望的，繞過了原有的資源分配邏輯。一種解決方式是在 YAML 文件中顯式設置容器的環境變數 NVIDIA_VISIBLE_DEVICES，如下所示：
apiVersion: v1kind: Podmetadata: name: tensorflowspec: containers: - name: tensorflow args: ["sleep 1d"] command: ["/bin/sh", "-c"] image: tensorflow/tensorflow:latest-gpu env: - name: NVIDIA_VISIBLE_DEVICES
回顧與展望
本文開始提到 Kubernetes 1.8 之前版本的 GPU 管理比較簡單，這裡對 1.8 之前版本的實現邏輯簡單介紹一下，並與本文介紹的 1.8 版本及之後的實現做對比，方便讀者瞭解 Kubernetes GPU 資源管理和調度的演進歷史，並能對當前實現方式的特點有直觀的認識。
- 設備發現：1.8 之前版本是通過直接匹配計算節點上的設備文件 /dev/nvidia* 、/dev/nvidia-uvm 和 /dev/nvidiactl。這是一種折中做法，並不是真實的查詢設備信息，存在不可靠的問題。資源的名稱 http://alpha.kubernetes.io/nvidia-gpu ，和當前的 http://nvidia.com/gpu 也有區別。
- 設備分配：1.8 之前版本用戶需要在 YAML 文件中顯式掛載驅動相關動態庫，1.8 之後版本用戶不需要掛載。除此之外，在 Allocate 這部分，1.8 之前版本是通過 docker 的 API 實現，這其實不滿足軟體通用性的要求，也是一個臨時的折中做法。1.8 之後版本是通過 Device Manager + Device Plugin 實現。
- 代碼耦合：1.8 之前版本 GPU 資源管理的代碼耦合在 Kubernetes 代碼中，不利於社區貢獻，而且增加了 Kubernetes 穩定運行的風險。
- 基於 Device Manager + Device Plugin 方式管理擴展硬體設備，不侷限於 GPU，還可以管理 InfiniBand，高性能網卡，FPGA 等。
- 1.8 之後的版本新增了設備健康檢查，此功能依賴於具體設備的 Device Plugin 實現，1.8 之前版本無健康監控功能。
目前，Kubernetes 社區對 GPU 的功能實現也在快速迭代，大致集中在兩個方向：
- 當前 GPU 調度的數量均以整數為單位，考慮到利用 Kubernetes 做 GPU Inference，在後續功能改進上，將來有可能實現類似 0.5 這種非整數的調度單位，多容器之間共享 GPU 計算資源。
- 除了繼續完善 GPU 管理調度的功能外，Kubernetes 與常用 AI 計算框架的結合也是社區工作的重點，兩者的緊密配合將會帶來更加便捷和高效的 AI 計算。
總結
在利用容器技術簡化 AI 計算框架單機運行的基礎上，本文詳細介紹了利用 Kubernetes 實現集羣的 GPU 資源調度和管理，重點介紹了當前的 Device Plugin + Device Manager 實現邏輯。並以 Nvidia 官方的 k8s-device-plugin 為例，分析了 k8s-device-plugin 與 Device Manager 的功能交互過程。在介紹 Device Manager 功能點之後，我們從用戶的角度梳理了在 Kubernetes 上 GPU 任務提交的工作流程，並針對不同的應用場景，給出了三個示例。最後，我們回顧了 Kubernetes GPU 功能的演進歷史，對比了 1.8 前後兩個版本實現的優缺點，並展望了未來的發展方向。
參考
- https://github.com/kubernetes/kubernetes/tree/master/pkg/kubelet/cm/devicemanager
- https://github.com/kubernetes/community/blob/master/contributors/design-proposals/resource-management/device-plugin.md
- https://kubernetes.io/docs/tasks/manage-gpus/scheduling-gpus/
推薦閱讀：
