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


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.

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


免責聲明!

本站轉載的文章為個人學習借鑒使用,本站對版權不負任何法律責任。如果侵犯了您的隱私權益,請聯系本站郵箱yoyou2525@163.com刪除。



 
粵ICP備18138465號   © 2018-2025 CODEPRJ.COM