Ray 在微信 AI 計算中的大規模實踐
一、背景
微信現在已經成為人們日常生活中非常重要的組成部分,而隨著人工智能的發展,微信內也為用戶提供了多種涉及 AI 計算的服務體驗。例如,語音消息的文字轉換、視頻號的 AIGC 和推薦、掃一掃功能的圖像識別等。這些功能由于微信的用戶規模巨大,所以 AI 計算的服務規模也非常大。
為了應對大規模的 AI 計算任務,我們建立了 Astra 平臺,目前,有許多業務已經接入了 Astra 平臺,這些業務在平臺上構建了眾多 AI 算法服務,覆蓋了 LLM/多媒體處理等多個方面。
目前我們對于 Ray 的使用場景主要是 Ray Serve,在 Astra 平臺之前,我們團隊是一個純粹的后臺開發團隊。因此在我們在實際工作中,會更加深入的思考 AI 算法服務與傳統微服務之間的區別。
首先,關于應用規模,傳統的微服務一般最多只有幾千個節點、十來萬核。然而,在AI算法這種計算密集型的任務上,我們的 AI 算法服務往往需要數十萬節點,其計算資源需求可達數百萬核。這種超級應用對我們的模塊管理系統以及 K8S 集群提出了極高的要求,要支撐如此大規模的應用部署是非常困難的。
其次,隨著資源數量和資源種類的增加,部署復雜度快速升高。AI 算法服務對 GPU 資源有特殊需求,市場上存在多種類型的 GPU,例如 NVIDIA、紫霄、昇騰等品牌。這些不同型號的 GPU 需要特定的適配工作,這無疑增加了每次部署時的復雜性和工作量。
再者,運維難度亦不容忽視。微服務主要涉及業務邏輯,不存在復用性,但 AI 算法是一種純算法服務,并不存在業務邏輯,不同的業務可能都會需要分別部署獨立的集群,這也相應地提高了運維工作的復雜程度。
最后,成本問題同樣突出。當前 GPU 的價格高昂,算力資源十分昂貴。降低 AI 推理的成本并提高資源利用率也是我們的重要目標之一。
鑒于上述考慮,我們選擇了 Ray。它可以提供統一的分布式平臺,整合多種計算模式,構建一個完整的生態系統。
早在 2022 年,我們就注意到 Ray 的優勢,當時對于 Ray 的理解尚淺,但觀察到國內外眾多企業,包括 ChatGPT 等先進應用的成功案例,我們決定加大對 Ray 的投入,利用其將單機應用輕松擴展至分布式環境的能力,簡化我們的開發流程,并實現更高效的資源管理。
接下來介紹 Astra 平臺與 Ray 結合的整體架構。在這一架構中,平臺管理的基本單位是基于 Ray 的應用,其底層由我們自主設計的聯邦集群架構支撐。該架構底層連接至公司內部多個 K8S 集群,允許我們的應用部署于不同的 K8S 集群之上。這些 K8S 節點是我們預先從 K8S 申請的資源,每個節點上都會集成我們 Starlink 集群管理的 Agent、網絡穿透 P2P 下載組件以及 TFCC AI 運行時。
架構圖的下層展示了 Starlink 集群管理模塊,它主要提供了以下幾項核心能力:服務發現、負載均衡、容災調度以及應用調度。面對的主要挑戰是支持多達數十萬個節點的大規模集群管理,這意味著我們需要具備處理百萬級 Pod 節點的能力,并確保高效的集群管理。此外,為了降低成本,考慮到騰訊內部可能存在資源利用不充分的情況,我們致力于優化資源配置,充分利用那些利用率較低或空閑的資源。最后,為了簡化應用部署流程,我們為算法開發人員提供了便捷的方式,使他們能夠迅速且輕松地部署應用程序。
通過這樣的架構設計,我們不僅能夠應對大規模集群管理帶來的技術挑戰,還能夠在成本控制和部署效率上取得顯著進步,從而更好地支持微信平臺上的 AI 計算需求。
二、百萬級節點的集群管理
接下來,將詳細探討在處理百萬級節點集群管理時所面臨的挑戰及解決方案。
調度架構通常可以分為三種類型:單層、兩層以及共享調度。每種架構都有其特點和適用場景。
單層調度:這種架構類似于 K8S,適用于規模較小的集群,一般為數千個節點。單層調度的一個主要限制是其較低的并發度,因為整個集群通常只依賴于單一的調度器進行資源分配。
兩層調度:此架構預先將資源池中的機器分配給上層調度器,例如 Spark 采用的就是這種模式。相比單層調度,兩層調度具有更高的并發度,支持更大規模的集群。
共享調度:這一概念源于谷歌的 Omega 論文。其核心特點是支持無限數量的調度器,每個應用都有自己的調度器,且所有調度器都能查看到整個資源池的狀態。我們實現了樂觀調度策略,即每個調度器假設沒有沖突地分配資源,并在檢測到沖突時進行調整。這種方法不僅能夠顯著提升調度規模和資源利用率,還能大幅提高調度并發度,非常適合處理數十萬節點的大規模集群需求。
對于需要管理幾十萬甚至更多節點的超級應用而言,選擇共享調度架構幾乎是必然的。它通過允許多個調度器同時工作,確保了高并發度和高效的資源利用,滿足了大規模集群管理的需求。我們的 Astra 平臺結合 Ray,正是采用了這樣的共享調度策略,以應對微信 AI 計算中遇到的巨大挑戰。
接下來,將深入介紹 Astra 自研調度架構的整體設計,我們從下往上介紹。
最底層是部署在不同 K8S 集群上的節點,這些資源是預先分配好的。例如,特定業務可以向 K8S 申請所需的資源,并將其集成到我們的平臺中使用。每個 K8S 集群包含各自的節點,這些節點會定期向 resource 集群發送心跳信號,以報告其狀態。
Resource 集群會聚合這些心跳信息,并通過廣播的方式實現輕量級的狀態同步。這一機制保證了即使單個集群僅由幾臺機器組成,也能夠支持多達百萬級別的節點管理,并維持一個在線節點列表,作為高性能資源管理的基礎。
Resource 是整個系統的核心部分,它不僅負責收集和同步節點狀態,還提供了高效的資源管理能力。Resource 集群能夠實時掌握所有節點的最新狀態,為上層調度提供準確的信息。
在 Resource 集群之上,是 APP 級別的獨立調度器。每個應用程序(APP)都有自己的調度器,這些調度器可以直接訪問整個集群中所有節點的狀態信息。每個調度器的主要職責包括:資源調度、負載均衡、容災調度。得益于共享調度的設計,調度器每分鐘可以處理數萬節點的調度任務,極大提升了系統的響應速度和效率。
最上層是用戶操作界面,用于執行諸如擴縮容等操作。這一層直接與最終用戶交互,提供了直觀的操作體驗,使得用戶可以輕松管理其應用程序和服務。
這套架構構成了我們進行百萬級別集群管理的基礎,通過分層設計和資源共享調度機制,實現了對大規模節點的有效管理,同時保證了高并發度和資源利用率,滿足了微信 AI 計算中的復雜需求。此架構不僅支持了數十萬節點的大規模集群管理,還顯著降低了運維復雜度和成本,提高了整體的服務質量和用戶體驗。
三、高效穩定地利用低優資源
下面探討如何高效且穩定地利用低優資源。騰訊內部低優資源的總量非常龐大,但與此同時也伴隨著諸多挑戰,“天下沒有免費的午餐”,盡管低優資源的成本較低,但也帶來了新的問題。
以下是從我們系統監控中截取的幾張圖,能說明我們所遇到的一些問題。
首先,節點處于亞健康狀態的比例較高,例如 CPU 時間片可能會被搶占,或者節點可能隨時被高優先級任務搶占。在 Ray 的應用上,當 Ray 應用部署到低優資源時,節點被搶占會導致集群難以恢復。特別是在 Ray Cluster 錯誤處理方面,如果頭節點出現故障,會難以恢復。當前社區版僅提供通過 Redis 備份的方式進行恢復,但我們不希望在服務中引入額外的依賴。因此,我們希望有一種簡潔的解決方案,能快速的地解決這一問題。
針對亞健康節點及被移除節點的快速處理,我們采取了一系列措施以確保系統的高效與穩定。以下是具體方法:
首先,對于異常節點的下線,K8S 平臺通常不會提前很長時間通知即將縮容。因此,我們使用 K8S 的 Pre Stop 特性來迅速響應。由于我們的心跳檢測頻率為每秒一次,一旦觸發 Pre Stop,我們可以立即從在線列表中移除該節點。在3秒鐘內更新路由表,這意味著在 Pre Stop 觸發后的 4 秒內,該節點將被完全移除。即使節點隨后被資源平臺直接終止,也不會對我們的服務造成影響。
其次,對于性能較低的節點,我們會實時監控并收集其運行數據。通過自動調節權重的路由算法調整節點接收的請求量。具體而言,性能較差的節點將被賦予較低權重,而高性能節點則保持較高權重。這種機制確保了任務能夠根據節點的實際性能進行合理分配,從而優化了 serving 層的工作負載分配,優化了整體服務吞吐和耗時。
通過上述措施,我們不僅能夠快速處理異常節點,還能有效管理性能差異較大的節點,確保系統在面對低優資源固有問題時仍能保持高效穩定的運行。
從我們在 Astra 平臺上使用低優資源的效果圖中可以清晰地看到優化前后的對比。圖中的綠色線條代表優化之前的狀態,而紅色線條則展示了優化之后的結果。通過對比可以看出,優化后我們的失敗率顯著降低,同時平均處理時間也得到了大幅改善。
在低優資源上使用 Ray 時,我們必須良好的設計系統來容忍節點(包括 worker 和頭節點)被頻繁踢出的情況。為此,我們借鑒了 K8S 聯邦集群的概念,提出了 Ray 聯邦集群的架構。這一架構的核心特點是每個 Ray cluster 作為一個服務的最小單位,每個 Ray Cluster 都能夠獨立提供完整的服務。這種設計確保了即使任意一個 Ray cluster 出現故障,整體服務依然不受影響。
Ray 聯邦集群的特點包括:
- 服務單元化:每個 Ray cluster 都是一個獨立的服務單元,可以單獨提供完整的服務。
- 容災能力:通過增加更多的 Ray cluster 來實現容災。每個 cluster 都可以視為一個容災單元,從而增強了系統的魯棒性。
- 突破規模限制:由于我們的基礎設施本身支持百萬級節點管理,我們可以將每個 Ray cluster 設計得較小,從而突破單個 cluster 的規模限制。
- 無狀態服務般的容災:無論哪個節點或整個 cluster 掛掉,我們都可以迅速替換,確保服務的連續性。
具體容災方案如下:
- 頭節點故障處理:當頭節點出現心跳異常時,我們會立即從路由算法中移除該 cluster,并調度一個新的 Ray cluster 進行替換。
- Worker 節點故障處理:如果僅是 worker 節點故障,該 cluster 仍可繼續使用,只需替換故障的 worker 節點即可。
- 擴展性:我們的 Ray Cluster 機制支持縱向擴容,可以根據需要調整 worker 數量,從而解決了大規模部署下的擴展性問題。
通過上述措施,我們不僅提升了系統的容災能力,還確保了在大規模節點管理中的高效性和靈活性。此架構使得每個 Ray APP 可以擁有無數個請求者,集群能夠像無狀態服務一樣靈活應對節點或 cluster 級別的故障,確保了服務的高可用性和穩定性。
針對訓練任務的自動調速,我們開發了一套內部工具用于數據處理。鑒于低優資源隨時可能被移除,這會導致集群處理能力下降,我們對每個節點的性能進行了預估,即根據資源類型和數量來評估其處理能力。隨著集群資源指標的變化,我們會自動調整任務的消費速率,以確保服務的正常運行。自動調速機制如下:
- 性能預估:我們根據資源類型(如不同型號的 GPU)和可用資源的數量,動態評估每個節點的處理能力。這一預估考慮了多種卡型共存的情況,確保評估結果的準確性。
- 動態調整消費速率:基于性能預估的結果,當檢測到集群中的資源變化時,系統會自動調整任務的消費速率。這種調整確保即使在資源減少的情況下,也能維持服務的穩定性和性能。
- 保障服務連續性:通過上述機制,即使低優資源被移除或集群處理能力發生變化,我們的系統仍能保持正常運行,避免因資源波動導致的服務中斷。
通過這套自動調速機制,我們的離線任務可以穩定的在低優資源上進行部署,確保了在復雜多變的環境中,離線任務能夠持續高效地運行。
四、降低應用部署復雜度
接下來探討如何降低應用部署的復雜度。一個 APP 可能包含多個模型,每個模型需要部署在不同類型的 GPU 上,且每種 GPU 類型可能會超出系統單個集群的數量、規模限制。這種情況下,我們面臨的復雜度大約是 O(N^3)。為了有效降低這一復雜度,我們采取了以下措施:
首先,多模型擴展。我們選擇 Ray 作為解決方案的關鍵原因之一在于其能夠簡化多模型擴展。以往,服務中可能需要引入一個中間模塊來專門管理各個模型之間的調用。而現在使用 Ray,僅需一個腳本即可實現多模型的有效管理和調度,大大簡化了部署流程。
其次,多卡型擴展。針對不同類型的 GPU,我們設計了相應的擴展機制。這部分內容將在后續詳細說明。
最后,多模塊擴展。我們的架構還支持多模塊的擴展,確保系統能夠靈活應對不同類型和數量的組件需求。
通過上述措施,我們顯著降低了應用部署的復雜度,使得多模型、多卡型以及多模塊的管理變得更加高效和簡便。
接下來我們詳細描述我們如何做到這一點。首先,統一基座環境。
- 核心組件:該環境包含了所有必需的核心組件,例如 P2P 網絡穿透、TFCC 以及特定版本的 CUDA,以確保一致性和兼容性。
- 包管理:允許用戶使用 Conda 管理包,允許用戶選擇不同版本的 Python,并實現自定義環境配置。這一系統是我們自行開發的環境管理工具,旨在為用戶提供靈活的包管理能力。
- 自定義庫與編譯:支持用戶自定義使用各種庫及進行 CI(持續集成)編譯操作,以滿足多樣化的開發需求。
其次, TFCC 引擎的支持。
針對不同類型的 GPU(如 NVIDIA、昇騰等),特別是當 NVIDIA 卡用于 TensorRT 時可能需要不同的模型編譯,TFCC 引擎起到了關鍵作用,它能夠屏蔽底層 GPU 卡的具體類型,使得用戶無需關心底層硬件細節。用戶只需調用 TFCC 提供的推理接口,即可完成推理任務,無論底層使用何種 GPU 卡。
通過上述標準化的環境管理方案,不僅簡化了多依賴項的服務部署,還確保了跨多種 GPU 類型的無縫支持,極大提升了用戶體驗和服務靈活性。
最后是 P2P 下載優化。隨著 AI 模型規模的不斷擴大,正如前面提到當前的模型體積已經顯著增加。例如,一些常見的 72B 模型可能達到 140GB 的大小。將如此龐大的模型下發至各節點進行部署耗時會比較長。為此,我們使用 P2P 對下載能力進行了多項優化:
- NAT 穿透:實現了 NAT 穿透功能,確保在復雜網絡環境下也能高效傳輸數據。
- 節點限速與全局限速:為保證網絡資源的合理分配,引入了節點級別的限速和全局限速機制。
通過這些優化措施,我們現在能夠將一個大約 140GB 的模型在大約 10 分鐘左右的時間內完成下發。這大大縮短了模型部署的時間,提高了整體工作效率。
五、Astra-Ray 使用樣例
下面是在平臺上的一些截圖,幫助大家了解我們平臺的使用界面和部署流程。以下是具體步驟:
第一步是修改代碼,只需要按照 ray serve 做簡單的改寫即可。
第二步是將修改后的代碼打包并上傳至 Git 倉庫。在此過程中,需要選擇一個 Python 版本。平臺會根據選擇自動打包,生成一個可以在我們平臺上運行的部署包。
第三步是變更 Ray 配置,這實際上是調整服務的 deployment 配置。此步驟允許用戶通過配置變更來動態調整服務參數。考慮到在線服務的穩定性,我們設計了灰度發布機制。為了確保配置變更能夠平穩生效,而不對現有服務造成即時沖擊,用戶可以通過灰度發布的方式逐步應用新的配置。這種方式允許新配置在一部分流量中先行驗證,確保其穩定性和兼容性后再全面推廣。通過這一機制,用戶可以在不影響整體服務的情況下,安全地測試和應用新的配置調整,從而保障了服務的連續性和可靠性。
第四步是進行擴容。完成擴容后,可以立即看到服務啟動,并在平臺頁面上查看集群的機器列表。例如,展示的一個集群機器列表中,包含大約 4,000 多個請求者(requester),每個請求者對應 5 臺機器,總計超過 2 萬個節點。
此外,我們還提供了以下功能以支持更高效的管理和調試:
- Ray Dashboard:用戶可以通過 dashboard 查看請求的狀態和性能指標。
- 日志調試:平臺支持日志記錄和調試功能,幫助用戶快速定位和解決問題。
六、Q&A
Q1:關于低優資源的抖動狀態問題
問:您剛才提到的低優資源存在 3 秒更新路由表和每秒更新機制。對于一些狀態不穩定的資源,例如其性能可能一會兒好一會兒壞,這種情況是否會引發抖動?這種不穩定狀態是如何處理的?
答:理解您的問題,確實存在兩種情況:
性能波動:當節點的 CPU 被壓制時,可能會出現性能不穩定的情況。我們通過動態調整權重來應對這一問題,確保即使性能波動也不會影響服務的整體質量。
節點掉線與重啟:如果一個節點突然掉線后重新啟動,我們會將其視為一個新節點進行處理。若它快速重啟,可能不會特別關注;但如果觸發了 PreStop 機制,則該節點會被移除,并在重新上線時作為新節點加入資源池。
Q2:關于每秒更新頻率的問題
問:每秒更新的頻率是否過高?這是否會帶來額外的壓力?
答:每秒心跳確實會帶來較高的請求量,尤其在面對數十萬個節點時,會產生較高的 QPS 壓力。為此,我們采用了聚合擴散機制來減輕集群負擔。首先,我們要保證自身服務的性能;其次,在集群容量接近極限時,可以通過擴展更多集群來支持更大的規模。Starlink 架構并非局限于單一 Resource 集群,而是也可以靈活擴展至多個集群。
Q3:超級 APP 跨集群調度機制
問:超級 APP 如何實現跨多個 Ray Serve 的調度?是通過每個集群內部的調度器還是有更高層次的調度機制?
答:通過 RPC 的方式,不同 Ray 集群之間是不互通的,但是我們保證每個 Ray 集群都有完整的功能。
Q4:超級 APP 副本故障恢復機制
問:在一個超級 APP 擁有大量副本的情況下,如果部分副本掛掉,是由哪個調度器負責重啟這些副本?
答:由 APP 自身的調度器負責維護其所有 Ray Serve 的狀態。即便應用規模龐大,調度器也僅需定時掃描即可管理。假設一個超級 APP 有 100 萬個集群,掃描一次也只需一秒鐘左右的時間。因此,即使是超大規模的應用,其調度器也不太可能成為瓶頸。
Q5:聯邦集群的逐步擴展與負載均衡
問:聯邦集群是如何逐步擴展并淘汰舊集群的?特別是在進行版本升級時,如何保持負載均衡?
答:聯邦集群的擴展和淘汰是一個漸進的過程。頭節點并不執行具體的業務邏輯或 worker 任務,而是與 worker 等價,以減少對系統的影響。在做路由算法時,我們將所有節點視為等同,簡化了管理和調度過程。
Q6:在線服務 Runtime 優化
問:是否在 Ray 的 runtime 層面進行了特定優化?
答:目前我們尚未對 Ray 框架本身進行修改。該項目從今年上半年開始引入 Ray,因此現階段的重點在于集成而非深度定制。
Q7:TFCC 引擎的透明性與模型轉換
問:實際使用者是否需要關注推理引擎的選擇?模型轉換和優化是如何處理的?
答:用戶無需關注底層使用的具體推理引擎。我們在打包過程中將運行內容分為三個部分:軟件包、環境包和數據集(通常是模型)。上傳時,我們會對模型進行預處理,轉換成適用于各種 GPU 卡的版本,并對其進行性能優化。這一切都是自動完成的,用戶只需提供輸入參數。
Q8:推理引擎通用接口
問:不同推理引擎的特性(如 Dynamic Batch、Attention 機制)是否通過通用接口封裝?
答:當前并不是所有推理任務都使用 TFCC,但我們計劃未來增加支持。最常用的接口只需要提供輸入形狀等基本信息。對于不同推理引擎的特性,我們正在考慮如何通過通用接口進行適配,但目前尚未完全封裝。
Q9:Ray 聯邦集群間的通信
問:騰訊內部的 Ray 聯邦集群是否實現了分布式共享內存的打通?不同 Ray serve 之間是否可以直接通信?
答:不同 requester 之間的對象無法直接獲取。要訪問其他 requester 的對象,只能通過 API 接口進行間接訪問。
Q10:TFCC 與顯存管理
問:TFCC 是否實現了顯存的分布式管理?
答:目前還未實現。