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

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

Kubernetes支持通過Device-Plugin的方式來(lái)增加對(duì)默認(rèn)資源(CPU,Memory等)之外的設(shè)備支持，而第三方可以通過編寫相應(yīng)的Device-Plugin來(lái)增加對(duì)設(shè)備的支持。目前Nvidia也是通過這樣的方式對(duì)GPU進(jìn)行支持。

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

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

Nvidia-Docker是Nvidia官方對(duì)容器做的拓展，用以讓容器能夠支持Nvidia的GPU設(shè)備。據(jù)官方統(tǒng)計(jì)的數(shù)據(jù)標(biāo)明，目前Nvidia-Docker的下載量已經(jīng)超過200萬(wàn)次，可以得知目前使用Nvidia-Docker來(lái)做AI系統(tǒng)環(huán)境已經(jīng)是非常主流的做法。

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

（圖片來(lái)源：https://devblogs.nvidia.com/gpu-containers-runtime）

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

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

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

如果在Docker配置中已經(jīng)做過相關(guān)配置，那么就可以簡(jiǎn)化為：

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

這里NVIDIA_VISIBLE_DEVICES這個(gè)環(huán)境變量用來(lái)指定需要綁定的GPU卡的邏輯ID，就可以實(shí)現(xiàn)容器中對(duì)該卡的綁定，使用上非常簡(jiǎn)單。

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

Device-Plugin機(jī)制本質(zhì)上是一個(gè)RPC服務(wù)。K8s定義了一個(gè)RPC調(diào)用接口，第三方資源設(shè)備方可以通過實(shí)現(xiàn)該接口來(lái)讓該設(shè)備得以在K8s側(cè)得以支持，并且在使用方式上和默認(rèn)資源沒有太大區(qū)別。Device-Plugin以Daemonset的方式在主機(jī)側(cè)運(yùn)行，并且通過一個(gè)Socket文件與Kubelet進(jìn)行通信，從而通過Kubelet給K8s上報(bào)相關(guān)信息。部署了該Daemonset的主機(jī)節(jié)點(diǎn)在k8s看來(lái)會(huì)包含由Device-Plugin注冊(cè)的硬件資源。Device-Plugin總的原理如下：

（圖片來(lái)源：https://medium.com/@Alibaba_Cloud）

首先，Device-Plugin需要向K8s注冊(cè)該資源，注冊(cè)機(jī)制通過實(shí)現(xiàn)以下RPC接口得以實(shí)現(xiàn)：

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

在詳細(xì)的rpc調(diào)用中，該接口會(huì)上報(bào)socket名稱、Device-Plugin的Api Version等信息，當(dāng)然，更重要的是它會(huì)上報(bào)ResourceName,該ResourceName會(huì)被K8s登記為該自定義設(shè)備的名稱，而名稱的規(guī)范是vendor-domain/resource，例如，Nvidia的GPU就被定義為nvidia.com/gpu，在用戶申請(qǐng)資源時(shí)，就需要使用該名稱。例如，在創(chuàng)建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

注冊(cè)之后，Device-Plugin還需要上報(bào)主機(jī)側(cè)的設(shè)備的數(shù)量和狀態(tài)，例如，如果主機(jī)節(jié)點(diǎn)上有8塊GPU卡，Device-Plugin會(huì)將該數(shù)量的資源數(shù)量和資源id列表告知K8s。此外，當(dāng)有Pod向K8s申請(qǐng)?jiān)撡Y源時(shí)，K8s會(huì)從上報(bào)的id列表中按照一定策略返回滿足用戶需求的資源數(shù)量的id序列，當(dāng)該id列表返回給Device-Plugin，再由Device-Plugin根據(jù)一定策略映射到真正的資源設(shè)備。以上的過程主要由以下的RPC調(diào)用實(shí)現(xiàn)的：

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

這里的ListAndWatch將由Device-Plugin調(diào)用NVML庫(kù)獲取主機(jī)側(cè)的GPU設(shè)備和狀態(tài)，并返回給k8s相應(yīng)的設(shè)備列表。而Allocate將在容器創(chuàng)建被調(diào)用，用來(lái)返回一些能夠使用主機(jī)上該資源的特殊配置，比如一些環(huán)境變量，再將這些信息給到Kubelet，并在容器啟動(dòng)的時(shí)候傳給容器。對(duì)于Nvidia GPU而言，主要是前面的提到的環(huán)境變量NVIDIA_VISIBLE_DEVICES，這樣容器啟動(dòng)的時(shí)候就能夠掛載相應(yīng)的GPU設(shè)備了。

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

這樣的方式在某些場(chǎng)景下是合理的，比如在強(qiáng)計(jì)算量任務(wù)的模式下，能夠避免不同進(jìn)程對(duì)GPU卡資源的爭(zhēng)搶以致發(fā)生顯存OOM等現(xiàn)象。但是在某些場(chǎng)景下缺會(huì)造成巨大的資源浪費(fèi)。比如有某些容器環(huán)境僅僅是給到用戶進(jìn)行算法的開發(fā)和調(diào)試，這樣的任務(wù)在絕大部分時(shí)間里并不會(huì)實(shí)際使用GPU，這種情況下能夠讓容器可以適當(dāng)?shù)墓蚕鞧PU是有很大價(jià)值的，畢竟GPU的價(jià)格非常昂貴，我們需要提供共享機(jī)制最大化資源的使用。此外Nvidia-Device-Plugin并沒有考慮GPU的親和性，這有可能會(huì)讓單容器多卡的容器遭遇較差的計(jì)算性能。這里會(huì)介紹我們的實(shí)現(xiàn)思路。

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

解決這個(gè)問題我們?cè)u(píng)估過兩個(gè)方案：調(diào)用NVML獲取GPU卡的實(shí)時(shí)狀態(tài)，并選擇負(fù)載較少的卡進(jìn)行分配；借助外部數(shù)據(jù)庫(kù)存放每個(gè)節(jié)點(diǎn)的GPU負(fù)載信息，Device-Plugin在Allocate的時(shí)候調(diào)用redis的信息查看負(fù)載情況。我們最終采用的是第二種方法，因?yàn)镹VML只能夠查詢進(jìn)程占用、資源占用的情況，如果一個(gè)容器綁定了一塊GPU卡，但容器里面并沒有任何進(jìn)程使用GPU，那么用NVML并不能查看真實(shí)的容器綁定關(guān)系。我們現(xiàn)階段還是希望把負(fù)載局限在卡上綁定的容器數(shù)，而不是真實(shí)的GPU使用率。

當(dāng)決定采用借助于一個(gè)外部的Redis數(shù)據(jù)庫(kù)的方案，用來(lái)存放每個(gè)節(jié)點(diǎn)的實(shí)時(shí)動(dòng)態(tài)的狀態(tài)，我們會(huì)在redis中為每一個(gè)節(jié)點(diǎn)維護(hù)一個(gè)單獨(dú)的map，記錄每個(gè)GPU id上分配的容器數(shù)量和任務(wù)名稱。當(dāng)Device-Plugin每次決定Allocate分配時(shí)，會(huì)去查詢Redis中查詢?cè)摴?jié)點(diǎn)下各個(gè)GPU id的容器負(fù)載，選擇最低的部分進(jìn)行綁定，并且將相應(yīng)的GPU id上的容器負(fù)載增加1。當(dāng)資源釋放時(shí)，Device-Plugin并不能知道釋放的消息，我們通過K8s的調(diào)度服務(wù)的Informer機(jī)制，在自定義的Informer中捕捉到釋放的POD的節(jié)點(diǎn)信息和任務(wù)名稱，并以此并將Redis中相應(yīng)的GPU id的資源數(shù)減去1。通過這種方式維護(hù)和監(jiān)控著GPU資源的分配信息。這里僅僅介紹解決的思路，具體細(xì)節(jié)不再展開。

GPU的通道親和性在多卡訓(xùn)練任務(wù)中非常重要，因?yàn)椴煌倪B接介質(zhì)對(duì)卡與卡之間的數(shù)據(jù)傳輸速度影響非常大。以下是一個(gè)典型的8卡GPU的卡間通道拓?fù)鋱D?？梢钥闯鲇械目ㄖg是通過Nvlink(NV1等)相連，有的是通過PCIe(PIX)相連。

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

（圖片來(lái)源：https://www.nvidia.com）

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

如果不考慮真實(shí)的GPU任務(wù)負(fù)載，單純的讓高親和性的卡綁定到一起是比較容易實(shí)現(xiàn)的。類似于共享GPU實(shí)現(xiàn)的思路，在多卡綁定的任務(wù)重，我們可以在Device-Plugin里面調(diào)用NVML，獲得GPU卡的連接拓?fù)?，從而知道它們之間的親和性關(guān)系。然后當(dāng)Allocate的時(shí)候，選擇讓高親和性通道間的GPU卡分配到一起即可。但是，如果考慮到高親和性的卡中間有部分的容器任務(wù)負(fù)載非常高，那么這個(gè)時(shí)候可能要高負(fù)載的影響。比較合理的方式是使用評(píng)分函數(shù)，按照一定策略給不同的可選組合評(píng)分，選擇得分最高的組合。但是我們采用較簡(jiǎn)單和直接的策略：首先選出負(fù)載最小的一個(gè)GPU id，再選擇跟這個(gè)id在同一高親和性通道的GPU卡，并挑選其中任務(wù)負(fù)載最小的卡進(jìn)行綁定。具體的細(xì)節(jié)不再展開了，主要是對(duì)NVML庫(kù)的調(diào)用，可以拿到主機(jī)的通道拓?fù)洌Ｏ碌墓ぷ魇琼樌沓烧碌摹?/p>

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