一文了解 DataLeap 中的 Notebook
精選一、概述
Notebook 是一種支持 REPL 模式的開發環境。所謂「REPL」,即「讀取-求值-輸出」循環:輸入一段代碼,立刻得到相應的結果,并繼續等待下一次輸入。它通常使得探索性的開發和調試更加便捷。在 Notebook 環境,你可以交互式地在其中編寫你的代碼、運行代碼、查看輸出、可視化數據并查看結果,使用起來非常靈活。
在數據開發領域,Notebook 廣泛應用于數據清理和轉換、數值模擬、統計建模、數據可視化、構建和訓練機器學習模型等方面。
但是顯然,做數據開發,只有 Notebook 是不夠的。在火山引擎 DataLeap 數據研發平臺,我們提供了任務開發、發布調度、監控運維等一系列能力。我們將 Notebook 作為一種任務類型,加入了數據研發平臺,使用戶既能擁有 Notebook 交互式的開發體驗,又能享受一站式大數據研發治理套件提供的便利。如果還不夠直觀的話,試想以下場景:
在交互式運行和可視化圖表的加持下,你很快就調試完成了一份 Notebook。簡單整理了下代碼,根據使用到的數據配置了上游任務依賴,上線了周期調度,并順手掛了報警。之后,基本上就不用管這個任務了:不需要每天手動檢查上游數據是否就緒;不需要每天來點擊運行,因為調度系統會自動幫你執行這個 Notebook;執行失敗了有報警,可以直接上平臺來處理;上游數據出錯了,可以請他們發起深度回溯,統一修數。
二、選型
2019 年末,在決定要支持 Notebook 任務的時候,我們調研了許多 Notebook 的實現,包括 Jupyter、Polynote、Zeppelin、Deepnote 等。Jupyter Notebook 是 Notebook 的傳統實現,它有著極其豐富的生態以及龐大的用戶群體,相信許多人都用過這個軟件。事實上,在字節跳動數據平臺發展早期,就有了在物理機集群上統一部署的 Jupyter(基于多用戶方案 JupyterHub),供內部的用戶使用。考慮到用戶習慣和其強大的生態,Jupyter 最終成為了我們的選擇。
Jupyter Notebook 是一個 Web 應用。通常認為其有兩個核心的概念:Notebook 和 Kernel。
Notebook 指的是代碼文件,一般在文件系統中存儲,后綴名為ipynb。Jupyter Notebook 后端提供了管理這些文件的能力,用戶可以通過 Jupyter Notebook 的頁面創建、打開、編輯、保存 Notebook。在 Notebook 中,用戶以一個一個 Cell 的形式編寫代碼,并按 Cell 運行代碼。Notebook 文件的具體內容格式,可參考 The Notebook file format (https://nbformat.readthedocs.io/en/latest/format_description.html)。
Kernel 是 Notebook 中的代碼實際的運行環境,它是一個獨立的進程。每一次「運行」動作,產生的效果是單個 Cell 的代碼被運行。具體來講,「運行」就是把 Cell 內的代碼片段,通過 Jupyter Notebook 后端以特定格式發送給 Kernel 進程,再從 Kernel 接受特定格式的返回,并反饋到頁面上。這里所說的「特定格式」,可參考 Messaging in Jupyter(https://jupyter-client.readthedocs.io/en/stable/messaging.html)。
在 DataLeap 數據研發平臺,開發過程圍繞的核心是任務。用戶可以在項目下的任務開發目錄創建子目錄和任務,像 IDE 一樣通過目錄樹管理其任務。Notebook 也是一種任務類型,用戶可以啟動一個獨立的任務 Kernel 環境,像開發其他普通任務一樣使用 Notebook。
三、技術路線
在 Jupyter 的生態下,除了 Notebook 本身,我們還注意到了很多其他組件。彼時,JupyterLab 正在逐漸取代傳統的 Jupyter Notebook 界面,成為新的標準。JupyterHub 使用廣泛,是多用戶 Notebook 的版本答案。脫胎于 Jupyter Kernel Gateway(JKG)的 Enterprise Gateway(EG),提供了我們需要的 Remote Kernel(上述的獨立任務 Kernel 環境)能力。2020 上半年,我們基于上面的三大組件,進行二次開發,在字節跳動數據研發平臺發布了 Notebook 任務類型。整體架構預覽如圖。
JupyterLab
前端這一側,我們選擇了基于更現代化的 JupyterLab (https://jupyterlab.readthedocs.io/en/stable/getting_started/overview.html) 進行改造。我們刨去了它的周邊視圖,只留下了中間的 Cell 編輯區,嵌入了 DataLeap 數據研發的頁面中。為了和 DataLeap 的視覺風格更契合,從 2020 下半年到 2021 年初,我們還針對性地改進了 JupyterLab 的 UI。這其中包括將整個 JupyterLab 使用的代碼編輯器從 CodeMirror 統一到 DataLeap 數據研發使用的 Monaco Editor,同時還接入了 DataLeap 提供的 Python & SQL 代碼智能補全功能。
額外地,我們還開發了定制的可視化 SDK,使得用戶在 Notebook 上計算得到的 Pandas Dataframe 可以接入 DataLeap 數據研發已經提供的數據結果分析模塊,直接在 Notebook 內部做一些簡單的數據探查。
JupyterHub
JupyterHub(https://jupyterhub.readthedocs.io/en/stable/) 提供了可擴展的認證鑒權能力和環境創建能力。首先,由于用戶較多,因此為每個用戶提供單獨的 Notebook 實例不太現實。因此我們決定,按 DataLeap 項目來切分 Notebook 實例,同項目下的用戶共享一個實例(即一個項目實際上在 JupyterHub 是一個用戶)。這也與 DataLeap 的項目權限體系保持了一致。注意這里的「Notebook 實例」,在我們的配置下,是拉起一個運行 JupyterLab 的環境。另外,由于我們會使用 Remote Kernel,所以在這個環境內,并不提供 Kernel 運行的能力。
在認證鑒權方面,我們讓 JupyterHub 請求我們業務后端提供的驗證接口,判斷登錄態的用戶是否具備請求的對應 DataLeap 項目的權限,以實現權限體系對接。
在環境創建方面,我們通過 OpenAPI 對接了字節跳動內部的 PaaS 服務,為每一個使用了 Notebook 任務的 DataLeap 項目分配一個 JupyterLab 實例,對應一個 PaaS 服務。由于直接新建一個服務的流程較長,速度較慢,因此我們還額外做了池化,預先啟動一批服務,當有新項目的用戶登入時直接分配。
Enterprise Gateway
Jupyter Enterprise Gateway (https://jupyter-enterprise-gateway.readthedocs.io/en/latest/) 提供了在分布式集群(包括 YARN、Kubernetes 等)內部啟動 Kernel 的能力,并成為了 Notebook 到集群內 Kernel 的代理。在原生的 Notebook 體系下,Kernel 是 Jupyter Notebook / JupyterLab 中的一個本地進程;對于啟用了 Gateway 功能的 Notebook 實例,所有 Kernel 相關的功能的請求,如獲取 Kernel 類型、啟動 Kernel、運行 Cell、中斷等,都會被代理到指定的 Gateway 上,再由 Gateway 代理到具體集群內的 Kernel 里,形成了 Remote Kernel 的模式。
這樣帶來的好處是,Kernel 和 Notebook 分離,不會相互影響:例如某個 Kernel 運行占用物理內存超限,不會導致其他同時運行的 Kernel 掛掉,即使他們都通過同一個 Notebook 實例來使用。
EG 本身提供的 Kernel 類型,和字節跳動內部系統并不完全兼容,需要我們自行修改和添加。我們首先以 Spark Kernel 的形式對接了字節跳動內部的 YARN 集群。Kernel 以 PySpark 的形式在 Cluster 模式的 Spark Driver 運行,并提供一個默認的 Spark Session。用戶可以通過在 Driver 上的 Kernel,直接發起運行 Spark 相關代碼。同時,為了滿足 Spark 用戶的使用習慣,我們額外提供了在同一個 Kernel 內交叉運行 SQL 和 Scala 代碼的能力。
2020 下半年,伴隨著云原生的浪潮,我們還接入了字節跳動云原生 K8s 集群,為用戶提供了 Python on K8s 的 Kernel。我們還擴展了很多自定義的能力,例如支持自定義鏡像,以及針對于 Spark Kernel 的自定義 Spark 參數。
穩定性方面,在當時的版本,EG 存在異步不夠徹底的問題,在 YARN 場景下,單個 EG 進程甚至只能跑起來十幾個 Kernel。我們發現了這一問題,并完成了各處所需的 async 邏輯改造,保證了服務的并發能力。另外,我們利用了字節跳動內部的負載均衡(nginx 七層代理集群)能力,部署多個 EG 實例,并指定單個 JupyterLab 實例的流量總是打到同一個 EG 實例上,實現了基本的 HA。
四、架構升級
當使用 Notebook 的項目日漸增加時,我們發現,運行中的 PaaS 服務實在太多了,之前的架構造成了
部署麻煩。全量升級 JupyterLab 較為痛苦。盡管有升級腳本,但是通過 API 操作升級服務,可能由于鏡像構建失敗等原因,會造成卡單現象,因此每次全量升級后都是人工巡檢檢查升級狀態,卡住的升級單人工點擊下一步。同時由于升級不同服務不會復用配置相同的鏡像,所以有多少服務就要構建多少次鏡像,當服務數量達到一定量級時,我們的批量升級請求可能把內部鏡像構建服務壓垮。
JupyterLab 需要不斷的根據用戶增長(項目增長)進行擴容,一旦預先啟動好的資源池不夠,就會存在新項目里有用戶打開 Notebook,需要經歷整個 JupyterLab 服務創建、環境拉起的流程,速度較慢,影響體驗。而且,JupyterLab 數量巨大后,遇到 bad case 的幾率增高,有些問題不易復現、非常偶發,重啟/遷移即可解決,但是在遇到的時候,用戶體驗受影響較大。
運維困難。當用戶 JupyterLab 可能出現問題,為了找到對應的 JupyterLab,我們需要先根據項目對應到 JupyterHub user,然后根據 user 找到 JupyterHub 記錄的服務 id,再去 PaaS 平臺找服務,進 webshell。
當然,還有資源的浪費。雖然每個實例很小(1c1g),但是數量很多;有些項目并不總是在使用 Notebook,但 JupyterLab 依然運行。
穩定性存在問題。一方面,JupyterHub 是一個單點,升級需要先起后停,掛了有風險。另一方面,EG 入流量經過特定負載均衡策略,本身是為了使 JupyterLab 固定往一個 EG 請求。在 EG 升級時,JupyterLab 請求的終端會隨之改變,極端情況下有可能造成 Kernel 啟動多次的情況。
基于簡化運維成本、降低架構復雜性,以及提高用戶體驗的考慮,2021 上半年,我們對整體架構進行了一次改良。在新的架構中,我們主要做了以下改進,大致簡化為下圖:
- 移除 JupyterHub,將 JupyterLab 改為多實例無狀態常駐服務,并實現對接 DataLeap 的多用戶鑒權。
- 改造原本落在 JupyterLab 本地的數據存儲,包括用戶自定義配置、Session 維護和代碼文件讀寫。
- EG 支持持久化 Kernel,將 Kernel 遠程環境元信息持久化在遠端存儲(MySQL)上,使其重啟時可以重連,且 JupyterLab 可以知道某個 Kernel 需要通過哪個 EG 連接。
鑒權 & 安全
單用戶的 Jupyter Notebook / JupyterLab 的鑒權相對簡單(實際上 JupyterLab 直接復用了 Jupyter Notebook 的這套代碼)。例如,使用默認命令啟動時,會自動生成一個 token,同時自動拉起瀏覽器。有了 token,就可以任意地訪問這個 Notebook。
事實上,JupyterHub 也是起到了維護 token 的作用。前端會發起一個獲取 token 的 API 請求,再拿著獲取的 token 請求通過 JupyterHub proxy 到真實的 Notebook 實例。而我們直接為 Jupyter Notebook 增加了 Auth 的功能,實現了在 JupyterLab 單實例上完成這套鑒權(此時,使用了 DataLeap 服務簽發的 Token)。
最后,由于所有用戶會共享同一組 JupyterLab,我們還需要禁止一些接口的調用,以保證系統的安全。最典型的接口包括關閉服務(Shutdown),以及修改配置等。后續 Notebook 所需的配置,轉由前端保存在瀏覽器內。
代碼 & Session 持久化
Jupyter Notebook 使用 File Manager(https://github.com/jupyter-server/jupyter_server/blob/main/jupyter_server/services/contents/filemanager.py) 管理 Contents 相關讀寫(對我們而言主要是 Notebook 代碼文件),原生行為是將代碼存儲在本地,多個服務實例之間無法共享同一份代碼,而且遷移時可能造成代碼丟失。
為了避免代碼丟失,我們的做法是,把代碼按項目分別存儲在 OSS 上并直接讀寫,同時解決了一些由于代碼文件元信息丟失,并發編輯導致的其他問題。例如,當多個頁面訪問同一份代碼文件時,都會從 OSS 獲取最新的 code,當用戶存儲時,前端會獲取最新的代碼文件,比較該文件的修改時間同前端存儲的是否一致,如果不同,則說明有其它頁面存儲過,會提示用戶選擇覆蓋或是恢復。
Notebook 使用 Session 管理用戶到 Kernel 的連接,例如前端通過 POST /session 接口啟動 Kernel,GET /session 查看當前運行中的 Kernel。在 Session 處理方面,原生的 Notebook 使用了原生的 sqlite(in memory),見代碼(https://github.com/jupyter-server/jupyter_server/blob/main/jupyter_server/services/sessions/sessionmanager.py)。盡管我們并不明白這么做的意義何在(畢竟原生的 Notebook 重啟,一切都沒了),但我們順著這個原生的表結構繼續前進,引入了 sqlalchemy 對接多種數據庫,將 Session 數據搬到了 MySQL。
另一方面,由于我們啟動的 Kernel,有一部分涉及 Spark on YARN,啟動速度并不理想,因此早期我們增加了功能,若某個 path 已有正在啟動的 Kernel,則等其啟動完畢而不是再啟動一個新的。這個功能原先使用內存中的 set 實現,現在也移植到了數據庫上,通過 sqlalchemy 來訪問。
Kernel 持久化 & 訪問
在 Remote Kernel 的場景下,一個 JupyterLab 需要知道它的某個 Kernel 具體在哪個 EG 上。在之前一個項目一個 JupyterLab 的狀態下,我們通過負載均衡簡單處理這個問題:即一個 Server 總是只訪問同一個 Gateway。然而當 JupyterLab 成為無狀態服務時,用戶并非固定只訪問一個 JupyterLab,也就不能保證總訪問用戶 Kernel 所在的 EG。
另一個情況是,當 JupyterLab 或 EG 重啟時,其上的 Kernel 都會關閉。當我們升級相關服務時,總是需要通知用戶準備重啟 Kernel。因此,為了實現升級對用戶無感,我們在 EG 這層開發了持久化 Kernel 的特性。
Kernel Gateway 在啟動 Kernel 時,記錄了關于 Kernel 的一些元信息,包括啟動參數、連接 Kernel 使用的 IP/Port 等。有了這些信息,當一個 Kernel Gateway 重啟且 Remote Kernel 不關閉,就有辦法重新連接上。原本這些信息默認在內存 dict 中維護,開源倉庫中有一套存儲在本地文件的方案;基于這套方案,我們擴展了自研的存儲到 MySQL 的方案。
在多實例的場景下,每一個 EG 實例依然會接管的各自的一部分 Kernel,并記錄每個 Kernel 由誰接管(探活、Cull Idle、連接使用等)。在其關閉前,需要清除接管信息,以便下次啟動或其他實例啟動時撈起。
為了減少 client(正常是 JupyterLab) 任意訪問 EG 的情況,一方面我們沿用了負載均衡的策略,另一方面 JupyterLab 在請求 Kernel 相關操作前,會先請求 EG 一次,由 EG 決定 JupyterLab 具體請求哪一個 EG IP/Port。
當 EG 服務本身重啟或者升級時,會在進程退出之前去清除接管信息。當頁面繼續訪問時,JupyterLab 服務將會隨機分發相應請求,由其它的 EG 服務繼續接管。
收益
架構升級簡化后,整套 Notebook 服務的穩定性獲得了極大的提升。由于實現了用戶無感知的升級,不僅提升了用戶的使用體驗,運維的成本也同時降低了。
部署的成本也極大地降低,包括算力、人力的節省。由于剝離了內部依賴,我們得以將這套架構部署在各種公有云、私有化場景。
五、調度方案
在前面,我們重點關注了怎么將 Jupyter 這套應用嵌入到 DataLeap 數據研發中。這只覆蓋了我們 Notebook 任務的頁面調試功能。實際上,同時作為一個調度系統,我們還需要關心怎么調度一個 Notebook 任務。
首先,是和所有其他任務類型相同的部分:當 Notebook 任務所配置的上游依賴任務全部運行完畢,開始拉起本次 Notebook 任務的運行。我們會根據任務的版本創建一個任務的快照,我們稱之為任務實例,并將其提交到我們的執行器中。
對于 Notebook 任務,在實例運行前,我們會根據 Notebook 任務對應的版本,從 OSS 拷貝一份 Notebook 代碼文件,用于執行。在具體的執行流程中,我們使用了 Jupyter 生態中的 nbconvert (https://nbconvert.readthedocs.io/en/latest/) 來實現在沒有 Jupyter 應用的前提下在后臺運行這份 Notebook 文件,并將運行后得到的結果 Notebook 文件傳回 OSS。nbconvert 的工作原理比較簡單,且復用了 Jupyter 底層的代碼,具體如下:
- 根據指定的 Kernel Manager 或 Notebook 文件里的 Kernel 類型創建對應的 Kernel Manager(https://github.com/jupyter/jupyter_client/blob/main/jupyter_client/manager.py);
- Kernel Manger 創建 Kernel Client,并啟動一個 Kernel;
- 遍歷 Notebook 文件里的 Cell,調用 Kernel Client 執行 Cell 里的代碼;
- 獲取輸出結果,按照 nbformat 指定的 schema 填入 NotebookNode,并保存。
下圖是調度執行 Notebook 的 Kernel 運行流程和通過調試走 EG 的 Remote Kernel 運行流程對比。可以看出,它們的鏈路并沒有本質上的區別,只不過是在調度執行時,不需要交互式的 Kernel 通信,以及 EG 的這些 Kernel Launcher 使用了 embed_kernel 在同進程內啟動 Kernel 而已。走到最底層,它們都是使用了 ipykernel 的(其他語言 kernel 同理)。
六、未來工作
Notebook 任務已成為字節跳動內部使用較為高頻的任務類型。在火山引擎,我們也可以購買 DataLeap,即一站式大數據研發治理套件,開通交互式分析的版本,使用到 DataLeap 的 Notebook 任務。
有的時候,我們發現,我們有比 Jupyter 社區快半步的地方:比如基于 asyncio 異步優化的 EG;比如給 Notebook 增加 Auth 能力。但社區的發展也很快:比如社區將 Jupyter 后端相關的代碼實現,統一收斂到了jupyter_server;比如 EG 作者提出的 Kernel Provider 方案,令jupyter_server可以直接支持 Remote Kernel。
因此我們并未就此止步。目前,這套 Notebook 服務和 DataLeap 數據研發的其他前后端服務,仍存在著割裂。未來,我們希望精簡架構,實現徹底的整合,使 Notebook 并非以嵌入的形式融合在 DataLeap 的產品中,而是使其原生就在 DataLeap 數據研發中被支持,帶來更好的性能,同時又保留所有 Jupyter 生態帶來的強大功能。另一方面,隨著 DataLeap 數據研發平臺對流式數據開發的支持,我們也希望借助 Notebook 實現用戶對流式數據的探索、調試、可視化等功能的需求。相信不久的將來,Notebook 能夠實現流批一體化,來服務更加廣泛的用戶群體。