近年來AI技術的繁榮和深化，尤其是深度學習的崛起，離不開海量數據和計算力的提升。尤其是對Nvidia的GPU的利用，讓深度學習獲得幾十倍的性能提升，才徹底打開了AI想象空間。雖然智慧芯片近年來有著百花齊放的風景，最典型的例如Google的TPU，但是平心而論，從普惠意義和生態上，Nvidia的GPU仍然占有主導位置。

不過，Nvidia的GPU無疑是昂貴的，所以如何最大化利用好GPU的硬件資源，是每一個算力平臺型產品都要考慮的問題。比如，有多個用戶使用GPU服務器進行訓練時，如何保證資源合理的分配非常重要。得益于Nvidia公司為Docker寫的Runtime，也就是Nvidia-Docker，使得在Docker里使用GPU成為可能。從容器粒度來管理和使用GPU要比從主機角度容易很多，因為運行GPU的AI任務通常配置非常復雜，這種復雜包括管理員從管理GPU卡的分配和使用者切換不同的訓練環境，而容器可以封裝不同的訓練環境，很大程度上降低復雜性。此外，借助Kubernetes來管理Nvidia-Docker，使得GPU任務的分配更加簡單和合理，目前已成為幾乎所有主流的AI算力平臺的方案。

Kubernetes支持通過Device-Plugin的方式來增加對默認資源(CPU,Memory等)之外的設備支持，而第三方可以通過編寫相應的Device-Plugin來增加對設備的支持。目前Nvidia也是通過這樣的方式對GPU進行支持。

K8s + Nvidia-Device-Plugin的方式兩個限制：每一塊GPU同時最多只能被一個容器使用；沒有考慮GPU卡之間的通道親和性。這樣的方式已經滿足了部分AI算力需求場景，但是某些場景下還是有缺點和限制的: 每一塊GPU同時最多只能被一個容器使用，這在訓練模式下沒有任何問題，但是在開發調試模式下會造成巨大的資源浪費。因為開發調試模式下用戶大部分時間并沒有實質運行GPU的資源，但卻排他的獨占了昂貴的GPU。此外在多卡主機架構下，GPU卡直接的連接通常是不一樣的，有的通過Nvlink相連，有的通過PCIe，而不同的連接方式性能差別非常大。而沒有考慮同個主機里GPU卡直接的通道親和性時，也會給多卡計算時需要發生數據傳輸時(如all_reduce操作)帶來過高的通信開銷。那么自然而然，在同一個容器掛載多張GPU卡時，我們當然更希望會掛載通道親和性更好的卡。

本文會介紹K8s進行GPU調度的通用流程和我們的一些改造方案。包括用于支持容器GPU掛載的Nvidia-Docker、K8s中將GPU作為拓展資源調度的Device-Plugin機制，以及針對原生Nvidia-Device-Plugin存在的問題的改造方案。

Nvidia-Docker是Nvidia官方對容器做的拓展，用以讓容器能夠支持Nvidia的GPU設備。據官方統計的數據標明，目前Nvidia-Docker的下載量已經超過200萬次，可以得知目前使用Nvidia-Docker來做AI系統環境已經是非常主流的做法。

這里不詳細介紹Nvidia-Docker的原理了，詳細的原理可以閱讀其官方的設計文檔，這里只作簡單的介紹。從2015年開始，Docker容器的誕生了一套容器運行時標準OCI(Open Containers Initiative)，它包含容器運行時標準(Runtime-Spec)和 容器鏡像標準(Image-Spec)。而著名的Runc則是這套標準的一個默認實現，然而任何滿足該標準的實現都可以注冊為容器的運行時拓展。Containerd則包裝了Runc和其它功能如生命周期管理等，以Daemon的形式運行在主機。Nvidia-Docker正是基于這一套拓展標準增加Nvidia GPU的容器支持。Nvidia-Docker主要原理是將對GPU的支持放入一個兼容OCI標準的運行時庫拓展libnvidia-container中，并在Runtime的API中進行調用，在libnvidia-container中通過共享和調用主機側的nvidia-driver實現對GPU的支持。在容器啟動時，Runc會調用一個叫做nvidia-container-runtime-hook的hook，這個hook會去檢查相應的環境是否具有GPU的支持和一些環境檢查，完成之后容器啟動，在運行時容器內進程也是通過libnvidia-container暴露的接口進行交互，從而實現容器對GPU的透傳和運行時支持。

（圖片來源：https://devblogs.nvidia.com/gpu-containers-runtime）

值得注意的是，Nvidia-Docker容器會使用主機側的Nvidia-Driver，再上層的軟件棧如cuda/cudnn，AI框架等，則在容器里面提供。此外，多個Nvidia-Docker可以掛載同一個GPU，只要通過環境變量指定就好，并沒有數量上的限制。

為了方便理解后面Device-Plugin的機制，這里簡單介紹一下Nvidia-Docker掛載不同GPU設備的方式。Nvidia-Docker的使用非常簡單，它通過指定一個環境變量來指定將要掛載的GPU設備，不過要在Docker的配置文件中指定Docker的Runtime為Nvidia-Docker，或者通過命令行顯式指定也可以：

nvidia-docker run -e NVIDIA_VISIBLE_DEVICES=0,1 --runtime=nvidia -it tensorflow/tensorflow-gpu:v1.13 bash

如果在Docker配置中已經做過相關配置，那么就可以簡化為：

docker run -e NVIDIA_VISIBLE_DEVICES=0,1 -it tensorflow/tensorflow-gpu:v1.13 bash

這里NVIDIA_VISIBLE_DEVICES這個環境變量用來指定需要綁定的GPU卡的邏輯ID，就可以實現容器中對該卡的綁定，使用上非常簡單。

K8s通過Device-Plugin的機制對非默認的資源設備進行支持，例如RDMA設備、AMD GPU等，當然也包括本文最關心的Nvidia GPU。通過編寫相應的Device-Plugin，第三方資源設備上可以在K8s中添加對相應設備的支持，讓用戶獲得和原生資源(CPU,Memory等)近乎一樣的使用體驗。

Device-Plugin機制本質上是一個RPC服務。K8s定義了一個RPC調用接口，第三方資源設備方可以通過實現該接口來讓該設備得以在K8s側得以支持，并且在使用方式上和默認資源沒有太大區別。Device-Plugin以Daemonset的方式在主機側運行，并且通過一個Socket文件與Kubelet進行通信，從而通過Kubelet給K8s上報相關信息。部署了該Daemonset的主機節點在k8s看來會包含由Device-Plugin注冊的硬件資源。Device-Plugin總的原理如下：

（圖片來源：https://medium.com/@Alibaba_Cloud）

首先，Device-Plugin需要向K8s注冊該資源，注冊機制通過實現以下RPC接口得以實現：

service Registration {  rpc Register(RegisterRequest) returns (Empty) {}}

在詳細的rpc調用中，該接口會上報socket名稱、Device-Plugin的Api Version等信息，當然，更重要的是它會上報ResourceName,該ResourceName會被K8s登記為該自定義設備的名稱，而名稱的規范是vendor-domain/resource，例如，Nvidia的GPU就被定義為nvidia.com/gpu，在用戶申請資源時，就需要使用該名稱。例如，在創建POD的資源配置里，需要這樣指定該資源：

apiVersion: v1kind: Podmetadata:  name: demo-podspec:  containers:    - name: demo-container-1      image: k8s.gcr.io/pause:2.0      resources:        limits:          nvidia.com/gpu: 2

注冊之后，Device-Plugin還需要上報主機側的設備的數量和狀態，例如，如果主機節點上有8塊GPU卡，Device-Plugin會將該數量的資源數量和資源id列表告知K8s。此外，當有Pod向K8s申請該資源時，K8s會從上報的id列表中按照一定策略返回滿足用戶需求的資源數量的id序列，當該id列表返回給Device-Plugin，再由Device-Plugin根據一定策略映射到真正的資源設備。以上的過程主要由以下的RPC調用實現的：

  service DevicePlugin {        rpc ListAndWatch(Empty) returns (stream ListAndWatchResponse) {}        rpc Allocate(AllocateRequest) returns (AllocateResponse) {}  }

這里的ListAndWatch將由Device-Plugin調用NVML庫獲取主機側的GPU設備和狀態，并返回給k8s相應的設備列表。而Allocate將在容器創建被調用，用來返回一些能夠使用主機上該資源的特殊配置，比如一些環境變量，再將這些信息給到Kubelet，并在容器啟動的時候傳給容器。對于Nvidia GPU而言，主要是前面的提到的環境變量NVIDIA_VISIBLE_DEVICES，這樣容器啟動的時候就能夠掛載相應的GPU設備了。

值得注意的是，容器側并沒有任何gpu虛擬化或者顯存分配的策略，所以Nvidia-Device-Plugin分配的粒度是單張卡，并且絕對的一一映射，即上報的GPU數量就是絕對的GPU數量，k8s側負責資源的分配，再由Device-Plugin受K8s的返回的需要掛載GPU數量和ID，并將容器映射到實際的GPU設備上。

這樣的方式在某些場景下是合理的，比如在強計算量任務的模式下，能夠避免不同進程對GPU卡資源的爭搶以致發生顯存OOM等現象。但是在某些場景下缺會造成巨大的資源浪費。比如有某些容器環境僅僅是給到用戶進行算法的開發和調試，這樣的任務在絕大部分時間里并不會實際使用GPU，這種情況下能夠讓容器可以適當的共享GPU是有很大價值的，畢竟GPU的價格非常昂貴，我們需要提供共享機制最大化資源的使用。此外Nvidia-Device-Plugin并沒有考慮GPU的親和性，這有可能會讓單容器多卡的容器遭遇較差的計算性能。這里會介紹我們的實現思路。

前面介紹過，Device-Plugin通過ListAndWatch接口上報GPU資源列表，那么自然而然，我們會想，如果能夠偽造出更多虛擬的GPU ID給K8s，K8s在分配POD資源的時候返回虛擬的id，再由Device-Plugin映射回真實的id從而實現GPU卡的復用，就像虛擬內存地址到物理內存地址映射一樣，而虛擬內存可以比物理內存大很多。是的，主要的思路就是這樣，構造出虛擬的GPU設備id，在容器真正啟動的時候再映射到真實的GPU設備。但是有一個非常重要的問題需要解決：怎么樣保證GPU負載的大體上的平衡，而不會出現某些卡上綁定了太多的容器，而另一些卡上則沒有任何容器？這個現象是有可能出現的，容器的壽命不一，不斷的創建和刪除容器極有可能會造成GPU資源分配的不均勻。所以虛擬id到真實id的映射不能夠是通過一個簡單的線性映射關系來決定，而是需要通過考慮在任務創建時的真實的GPU負載來動態的決定掛載GPU的id。

解決這個問題我們評估過兩個方案：調用NVML獲取GPU卡的實時狀態，并選擇負載較少的卡進行分配；借助外部數據庫存放每個節點的GPU負載信息，Device-Plugin在Allocate的時候調用Redis的信息查看負載情況。我們最終采用的是第二種方法，因為NVML只能夠查詢進程占用、資源占用的情況，如果一個容器綁定了一塊GPU卡，但容器里面并沒有任何進程使用GPU，那么用NVML并不能查看真實的容器綁定關系。我們現階段還是希望把負載局限在卡上綁定的容器數，而不是真實的GPU使用率。

當決定采用借助于一個外部的Redis數據庫的方案，用來存放每個節點的實時動態的狀態，我們會在redis中為每一個節點維護一個單獨的map，記錄每個GPU id上分配的容器數量和任務名稱。當Device-Plugin每次決定Allocate分配時，會去查詢Redis中查詢該節點下各個GPU id的容器負載，選擇最低的部分進行綁定，并且將相應的GPU id上的容器負載增加1。當資源釋放時，Device-Plugin并不能知道釋放的消息，我們通過K8s的調度服務的Informer機制，在自定義的Informer中捕捉到釋放的POD的節點信息和任務名稱，并以此并將Redis中相應的GPU id的資源數減去1。通過這種方式維護和監控著GPU資源的分配信息。這里僅僅介紹解決的思路，具體細節不再展開。

GPU的通道親和性在多卡訓練任務中非常重要，因為不同的連接介質對卡與卡之間的數據傳輸速度影響非常大。以下是一個典型的8卡GPU的卡間通道拓撲圖。可以看出有的卡之間是通過Nvlink(NV1等)相連，有的是通過PCIe(PIX)相連。

而不同的GPU通道會導致完全不一樣的數據傳輸性能，通常它們之間的速度傳輸能相差很多倍，例如，Nvlink可以達到幾十GB/s，而PCIe通常只有10 GB/s左右的吞吐性能。下圖是Nvidia Tesla P100的系列的直觀的連通拓撲圖和通道傳輸性能：

（圖片來源：https://www.nvidia.com）

正如前面所說，Nvidia-Device-Plugin并沒有考慮通道的親和性，也就是說在一個單容器雙卡的容器中，通過K8s的調度極有可能將兩個通過PCIe連接的卡綁定到一起，即便有在同一個Nvlink通道的卡的存在，這顯然是不合理的。高親和性和容器負載均衡有時會是相互矛盾的需求，如果追求絕對的親和性，那么可能要犧牲容器任務的負載均衡。我們采用的算法策略是盡量讓這兩個指標取得一個均衡。

如果不考慮真實的GPU任務負載，單純的讓高親和性的卡綁定到一起是比較容易實現的。類似于共享GPU實現的思路，在多卡綁定的任務重，我們可以在Device-Plugin里面調用NVML，獲得GPU卡的連接拓撲，從而知道它們之間的親和性關系。然后當Allocate的時候，選擇讓高親和性通道間的GPU卡分配到一起即可。但是，如果考慮到高親和性的卡中間有部分的容器任務負載非常高，那么這個時候可能要高負載的影響。比較合理的方式是使用評分函數，按照一定策略給不同的可選組合評分，選擇得分最高的組合。但是我們采用較簡單和直接的策略：首先選出負載最小的一個GPU id，再選擇跟這個id在同一高親和性通道的GPU卡，并挑選其中任務負載最小的卡進行綁定。具體的細節不再展開了，主要是對NVML庫的調用，可以拿到主機的通道拓撲，剩下的工作是順理成章的。

本文簡單介紹了Docker對Nvidia GPU的支持方案以及K8s的Device-Plugin機制。并且針對現有的Nvidia-Device-Plugin的某些場景缺陷提供解決思路。主要是針對GPU卡的任務間共享和通道親和性的優化。然而，這些改造都是要看場景的，有的場景值得和迫切需要這么做，但是有的場景卻非常不適合。這樣的改造會增加外部依賴并且讓Nvidia-Device-Plugin的GPU卡的綁定策略變得更加復雜，所以我個人強烈建議只有在必要的時候進行這些改造。而平臺的交互式測試、驗證場景，正是這樣改造的場景和動力。