雲原生的彈性 AI 訓練系列之三:藉助彈性伸縮的 Jupyter Notebook,大幅提高 GPU 利用率

騰訊雲原生發表於2021-10-18

Jupyter Notebooks 在 Kubernetes 上部署往往需要繫結一張 GPU,而大多數時候 GPU 並沒有被使用,因此利用率低下。為了解決這一問題,我們開源了 elastic-jupyter-operator,將佔用 GPU 的 Kernel 元件單獨部署,在長期空閒的情況下自動回收,釋放佔用的 GPU。這篇文章主要介紹了這一開源專案的使用方式以及工作原理。

Jupyter Notebooks 是目前應用最為廣泛的互動式開發環境,它很好地滿足了資料科學、深度學習模型構建等場景的程式碼開發需求。不過 Jupyter Notebooks 在方便了演算法工程師和資料科學家們日常開發工作的同時,也對基礎架構提出了更多的挑戰。

資源利用率的問題

最大的挑戰來自於 GPU 資源利用率。在執行的過程中即使沒有程式碼在執行,Notebook 也會長期佔用著 GPU,造成 GPU 的空置等問題。在大規模部署 Jupyter 例項的場景下,一般會通過 Kubernetes 建立多個 Notebook 例項,分配給不同的演算法工程師使用。而在這樣的情況下,我們需要在對應的 Deployment 中事先申請 GPU,這樣 GPU 會與對應的 Notebook 例項繫結,每個 Notebook 例項都會佔用一張 GPU 顯示卡。

然而同一時間,並不是所有的演算法工程師都在使用 GPU。在 Jupyter 中,編輯程式碼的過程是不需要使用計算資源的,只有在執行 Cell 中的程式碼片段時,才會使用 CPU 或 GPU 等硬體資源,執行並返回結果。由此可以預見,如果通過這樣的部署方式會造成相當程度的資源浪費。

造成這一問題的原因主要是原生的 Jupyter Notebooks 沒有很好地適配 Kubernetes。在介紹問題原因之前,先簡單概述一下 Jupyter Notebook 的技術架構。如下圖所示,Jupyter Notebook 主要由三部分組成,分別是使用者和瀏覽器端,Notebook Server 和 Kernel。

其中使用者和瀏覽器端是 Jupyter 的前端,主要負責展示程式碼和執行結果等。Notebook Server 是它的後端伺服器,來自瀏覽器的程式碼執行請求會被 Notebook Server 處理,分派給 Kernel 執行。Kernel 是真正負責執行程式碼,返回結果。

在傳統的使用方式中,使用者會通過 jupyter notebook $CODE_PATH 等命令,在本地執行 Jupyter Notebook Server,隨後訪問瀏覽器中的 Jupyter 互動式開發介面。當程式碼需要執行時,Notebook Server 會建立一個獨立的 Kernel 程式,這一程式會使用 GPU 等執行。在 Kernel 長期空閒,沒有程式碼需要執行時,這一程式會被終止,GPU 也就不再會被佔用。

而當部署在 Kuberenetes 之上後,問題就產生了。Notebook Server 和 Kernel 執行在同一個 Pod 的同一個容器下,儘管只有執行程式碼時才需要執行的 Kernel 元件是需要 GPU 的,而長期執行的 Notebook Server 是不需要的,但是受限於 Kubernetes 的資源管理機制,還是需要給其提前申請 GPU 資源。

在 Notebook Server 的整個生命週期中,這一塊 GPU 始終與 Pod 繫結。在 Kernel 程式空閒時雖然會被回收,但是已經分配給 Pod 的 GPU 卡卻不能再交還給 Kubernetes 進行排程了。

解決方案

為了解決這一問題,我們開源了專案 elastic-jupyter-operator。思路非常樸素:問題源於 Notebook Server 和 Kernel 在同一個 Pod 中,導致我們無法分別為這兩個元件申請計算資源。那隻要將他們分開部署,讓 Notebook Server 在單獨的 Pod 中,Kernel 也在單獨的 Pod 中,相互之間通過 ZeroMQ 通訊即可。

通過這樣的方式,Kernel 會在空閒時被釋放。在需要時會再次被臨時性地申請 GPU,執行起來。為了實現這一目的,我們在 Kubernetes 中實現了 5 個 CRD,同時為 Jupyter 引入了一個新的 KernelLauncher 實現。通過它們,使用者可以在 GPU 空閒時將 Kernel 回收釋放,在需要執行程式碼時再動態地申請 GPU 資源,建立 Kernel Pod 進行程式碼執行。

簡單的例子

下面我們將通過一個例子介紹使用方式。首先我們需要建立 JupyterNotebook CR(CustomResource),這一個 CR 會建立出對應的 Notebook Server:

apiVersion: kubeflow.tkestack.io/v1alpha1
kind: JupyterNotebook
metadata:
  name: jupyternotebook-elastic
spec:
  gateway:
    name: jupytergateway-elastic
    namespace: default
  auth:
    mode: disable

其中指定了 gateway,這是另外一個 CR JupyterGateway。為了能夠讓 Jupyter 支援遠端的 Kernel,需要這樣一個閘道器進行請求的轉發。我們同樣需要建立這樣一個 CR:

apiVersion: kubeflow.tkestack.io/v1alpha1
kind: JupyterGateway
metadata:
  name: jupytergateway-elastic
spec:
  cullIdleTimeout: 3600
  image: ccr.ccs.tencentyun.com/kubeflow-oteam/enterprise-gateway:2.5.0

JupyterGateway CR 中的配置 cullIdleTimeout 指定了經過多久的空閒時間後,其管理的 Kernel Pod 會被系統回收釋放。在例子中是 1 個小時。建立完這兩個資源後,就可以體驗到彈性伸縮的 Jupyter Notebook 了。如果在一個小時內一直沒有使用的話,Kernel 會被回收。

$ kubectl apply -f ./examples/elastic/kubeflow.tkestack.io_v1alpha1_jupyternotebook.yaml
$ kubectl apply -f ./examples/elastic/kubeflow.tkestack.io_v1alpha1_jupytergateway.yaml
$ kubectl port-forward deploy/jupyternotebook-elastic 8888:8888
$ kubectl get pods 
NAME                                          READY   STATUS    RESTARTS   AGE
jovyan-219cfd49-89ad-428c-8e0d-3e61e15d79a7   1/1     Running   0          170m
jupytergateway-elastic-868d8f465c-8mg44       1/1     Running   0          3h
jupyternotebook-elastic-787d94bb4b-xdwnc      1/1     Running   0          3h10m

除此之外,由於 Notebook 和 Kernel 解耦的設計,使得使用者可以方便地修改 Kernel 的映象與資源配額、向已經在執行的 Notebook 中新增新的 Kernel 等。

設計與實現

在介紹完使用方式後,我們簡單介紹其設計與實現。

當使用者在瀏覽器中選擇執行程式碼時,首先請求會傳送給在 Kubernetes 上執行的 Notebook Server。由於目前叢集上沒有正在執行的 Kernel,程式碼執行任務無法分配下去,所以 Notebook Server 會向 Gateway 傳送一個建立 Kernel 的請求。Gateway 負責管理遠端的 Kernel 的生命週期,它會在 Kubernetes 叢集中建立對應的 JupyterKernel CR。隨後會與叢集中已經建立好的 Kernel 通過 ZeroMQ 進行互動,然後將程式碼執行的請求傳送給 Kernel 進行執行,隨後將結果傳送給 Notebook Server 再將其返回給前端進行渲染和展示。

而 Gateway 會根據在 JupyterGateway CR 中定義的有關資源回收的引數,定時檢查目前管理的 Kernel 中有沒有滿足要求,需要被回收的例項。當 Kernel 空閒時間達到了定義的閾值時,Gateway 會刪除對應的 JupyterKernel CR,將其回收,釋放 GPU。

總結

目前深度學習在開發與落地生產的過程中仍然存在著諸多的挑戰。elastic-jupyter-operator 瞄準了在開發過程中的 GPU 利用率與開發效率的問題,提出了一種可行的方案,將佔用 GPU 的 Kernel 元件單獨部署,在長期空閒的情況下自動回收,釋放佔用的 GPU,通過這樣的方式提高資源的利用率的同時,也給予了演算法工程師使用者更多的靈活度。

從演算法工程師的角度來說,elastic-jupyter-operator 支援自定義的 Kernel,可以自行選擇在 Kernel 的容器映象中安裝 Python 包或者系統依賴,不需要擔心與團隊內部的 Notebook 統一映象的版本一致性問題,提高研發效率。

而從運維與資源管理的角度來說,elastic-jupyter-operator 遵循了雲原生的設計理念,以 5 個 CRD 的方式對外提供服務,對於已經落地 Kuerbenetes 的團隊來說具有較低的運維成本。

License

  • This article is licensed under CC BY-NC-SA 3.0.
  • Please contact me for commercial use.

【騰訊雲原生】雲說新品、雲研新術、雲遊新活、雲賞資訊,掃碼關注同名公眾號,及時獲取更多幹貨!!

相關文章