本文由微信後臺Astra專案團隊分享，原題“Ray在微信AI計算中的大規模實踐”，下文進行了排版和內容最佳化。
1、引言
微信存在大量AI計算的應用場景，主要分為三種：流量分發、產品運營和內容創作。流量分發場景中的 AI 計算主要用於搜尋、廣告、推薦場景的核心特徵生產，產品運營相關的 AI 計算主要用於產品功能相關和內容運營相關（低質、優質、生態建設），由於大模型的興起，AIGC 相關的文生圖、圖生圖、AI 特效等內容創作場景的 AI 計算也有了較多的落地。目前AI 計算幾乎覆蓋了微信的所有業務場景。
▲ 圖 1：微信內 AI 計算應用場景
然而，我們在使用微信已有的後臺基礎設施實現AI應用時遇到各種問題：
1）在資源層面，AI應用屬於計算密集型，計算複雜度高，需要大量資源，直接使用線上資源會導致成本過高；
2）在部署層面，微信後臺常見的部署平臺更適合部署I/O密集、高併發、高請求量的微服務，而AI應用則需要適配大量異構硬體和異構資源平臺，部署複雜度呈指數級上升；
3）在應用編排層面，直接透過訊息佇列等基礎元件解決複雜特徵依賴及相關非同步過程，開發效率低，變更風險高，可觀測性差；
4）在平臺層面，由於缺乏平臺支撐，演算法迭代速度慢，模型能力使用門檻高。因此，微信亟需一個低成本、高效率、低門檻的AI計算平臺來解決上述問題。
▲ 圖 2：微信內原有基礎設施
比如，OCR作為影片號推薦和影片號搜尋依賴的一個重要特徵，計算量非常大，需要超過100 萬核的CPU計算資源，同時對實時性和可靠性的要求很高，需要在 1 分鐘內完成特徵生成。P6n平臺適合做高實時(毫秒級響應)的線上任務，實時性上可以滿足需求，但固定部署的資源成本較高，多模型部署複雜度高，不符合需求。Gemini 平臺更適合做大規模長時間的離線任務，在實時性和可靠性上不滿足需求。我們需要一個高實時（10 秒級響應），支援大規模異構資源部署，低成本和高可靠的近線任務平臺。
2、為何在AI計算中引入Ray？
▲ 圖 3：使用 Ray 構建 AI 計算的企業
Ray是一個通用的分散式計算引擎，2016年開源於加州大學伯克利分校 RISELab，是發展最快的計算引擎之一。目前已經廣泛應用於OpenAI、螞蟻、位元組和華為等公司，是新一代的明星計算框架。首先：編寫分散式計算既簡單又直觀。開發者不必瞭解所有通訊和排程細節，也不必對此進行推理。藉助 Ray 的簡單原語，可以將任何 Python 函式或類轉換為分散式執行：只需新增一個裝飾器，就大功告成了。Ray 的分散式API 很簡單，所有複雜性都由 Ray 的執行框架處理。函式將被安排為無狀態任務執行，而類將是一個有狀態的遠端服務。
def detect(image_data):
model = load_detect_model() return model(image_data)def recognize(det_result):
model = load_recognize_model() return model(det_result)def ocr(image_data):
det_result = detect(image_data) return recognize(det_result)image_data = load_image_data()
ocr_result = ocr(image_data)
以上是一個圖片ocr本地執行的 python 指令碼，如果使用微服務部署，因為模型過大，單機視訊記憶體不夠，無法載入所有模型，則需要部署三個微服務模組：detect、recognize和ocr，應用部署的複雜度較高。
@ray.remote(num_gpus=1,num_cpus=16)
def detect(image_data):
model = load_detect_model() return model(image_data)@ray.remote(num_gpus=2,num_cpus=16)
def recognize(detect_result):
model = load_recognize_model() return model(detect_result)@ray.remote(num_cpus=4)
def ocr(image_data):
det_result = detect.remote(image_data) return recognize.remote(det_result)image_data = load_image_data()
ocr_result = ocr.remote(image_data)
如果使用 ray 來做 ocr 推理，只需要新增裝飾器@remote，指定模型使用的 cpu 和 gpu 資源數，透過一個python 指令碼即可完成ocr應用的部署，效率提升至少一個數量級。
▲ 圖 4：Ray AIR 如何以簡單的方式統一 ML 庫
其次：大多數流行的 ML 庫都與 Ray 有很強的整合性，而且 Ray 的原生庫也整合了這些庫。例如，開發者可以輕鬆地將 XGBBoost 與 Ray Train 結合使用，可以輕鬆地將 HuggingFace 與 Ray Serve 結合使用。或者，可以輕鬆地將 PyTorch 和 TensorFlow 與 Ray Train 結合使用。簡而言之，它擁有豐富的整合生態系統，不僅與 ML 庫整合，還與其他工具和框架整合。
第三：開發人員可以使用膝上型電腦進行開發。當你想將其擴充套件到 Ray 叢集時，只需更改一行程式碼或不更改任何程式碼即可輕鬆完成。
RAY_ADDRESS=ray://<cluster>:<port> python your_script.py
總的來說，Ray提供了高效能的分散式框架和簡單的分散式原語，提供了統一的分散式底盤。Ray融合不同計算正規化，與眾多開源元件便捷地結合從而實現對現有流程的提效。同時，Ray有完善的生態，資料處理、訓練、推理和服務等AI基礎設施需要的主流框架都可以很方便地在Ray上進行整合，大量知名企業選用 Ray開發 AI 計算。綜上，我們選擇了Ray 作為微信 AI 計算平臺的分散式底座。
3、微信基於Ray的AstraRay平臺
P6n是基於 Kubernetes微服務部署平臺，透過自動化編排和彈性擴縮容機制，很好的解決了線上高實時的後臺服務運維自動化問題，但不支援大規模的批處理服務，單應用多模型的部署複雜度較高，機器成本較高，不適合“在離線一體”的 AI計算場景。Gemini 是基於 kubernetes 的大資料平臺，適合處理離線大規模的資料清洗和模型訓練，但是由於排程的實時性不夠，不適合高實時性、高吞吐的和高可靠的AI計算場景。Astra 平臺要實現高實時、高吞吐、高可靠、低成本的 AI 計算平臺，需要解決如下幾個核心問題。
比如：
1）為了低成本，需要支援各種異構資源擴充套件；
2）為了高吞吐，支援超大規模資源排程；
3）降低單應用多模型的部署複雜度。
我們基於 Ray 計算底座，解決了上述三個核心問題，構建出適合 AI 計算平臺：AstraRay，在微信內進行了大規模 AI 應用部署的實踐。AstraRay 相比社群版本Ray(KubeRay) 有以下改進：
4、AstraRay平臺架構概覽
▲ 圖 7：kuberay 架構
▲ 圖 8：KubeRay 提交任務流程
業界使用社群成熟的 KubeRay 方案，透過 Ray 和 K8s 結合，提供了易用、高可用、高伸縮的雲原生 Ray 叢集服務，可以滿足中小規模 AI 應用的需求。但它有叢集規模小（最大僅支援數千個節點），異構資源擴充套件困難（單個 ray 叢集只能部署在一個 k8s 叢集，不支援聯邦k8s 叢集）和伸縮慢（受限於 K8s 的擴縮容速度）的問題，不適合微信內超大規模 AI 應用的需求。
▲ 圖 9：AstraRay 整體架構
我們在落地 Ray 的過程中遇到了三個核心技術挑戰：
1）百萬級 pod 的叢集管理：在影片號業務場景中，有超過百萬核的超級應用，已經遠超 K8s 叢集上限，我們希望單個 Ray 應用能支援百萬級別的 pod 的擴充套件；
2）不穩定資源下構建穩定服務：由於 AI 計算的資源消耗大，為了降低成本，我們大量使用了低成本、閒置，但穩定性差的計算資源。我們希望可以在不穩定資源上提供可靠穩定的服務；
3）降低應用部署的複雜度：微信內 AI 應用遇到模型、硬體、模組三種維度的異構問題，部署複雜度高。
我們希望使用統一的應用維度來簡化應用部署，即將 O(n^3) 複雜度降低為 O(1)。Astra 的部署系統架構如上圖，在 Poseidon/算力/太極/Gemini 等多個資源平臺基礎上擴充套件多個tke模組，組成擁有數百萬核CPU、萬卡GPU級別的超大叢集。我們透過服務發現的架構設計，解決了百萬級pod叢集管理的問題，透過負載均衡和容災排程解決了不穩定資源構建穩定服務的挑戰，同時透過應用排程解決了多模型應用部署複雜度的問題。接下來詳細介紹我們如何應對這三個技術挑戰。
5、技術挑戰1：單叢集支援百萬級計算節點
5.1 架構選擇
▲ 圖 11：叢集排程架構分類
業界系統的排程架構主要分為四類：單體排程、兩層排程、共享排程和混合排程。這些排程架構的本質區別其實只有兩點：
1）排程時排程器是否擁有全域性的資源檢視；
2）不同的應用是否擁有多個資源排程器。
單體排程顧名思義，即只有一個排程器，排程器擁有全域性資源檢視的架構，Google Borg 和 K8s 都採用這個架構。單體架構的好處是，所有的任務都由唯一的排程器處理，排程器可以充分的考慮全域性的資源使用情況，能方便的做出最優排程。但由於排程架構的限制，叢集效能受限於單體的效能，無法支撐過大的叢集。兩層排程擁有多個排程器，Apache Mesos 和 Hadoop YARN 都採用這個架構。兩層排程中，每個應用的排程器首先向中心節點獲取資源，再將其分配給應用中的各個任務。兩層排程解決了單體排程的效能問題，但是排程器僅擁有區域性資源檢視，無法做出最優排程。共享排程擁有多個排程器，每個排程器擁有全域性資源檢視，Omega 採用了這個架構。共享排程方案中，每個排程器都可以併發地從整個資源池中申請資源，解決了效能問題和最優排程問題，且可以支援較大叢集。因此，AstraRay 選擇共享排程來支援超大規模的資源管理。排程器間資源申請衝突可透過悲觀鎖或樂觀鎖來解決，AstraRay 實現了基於樂觀鎖的方案，出現衝突後再處理，無需中心節點，併發度更高。
5.2 Starlink排程
我們提出了一個新的排程系統 Starlink 來更好適配異構資源和硬體。Starlink採用共享排程架構，透過樂觀併發排程處理衝突，支援部署在任何基礎資源平臺(K8s/Yard/CVM)之上，且允許單個應用執行於多種異構的資源節點上。
▲ 圖 12：Starlink 排程架構
Starlink主要分為四個部分：
1）Node：任意部署了 Starlink 的 Agent 節點都可以成為 Node，Node 每秒會向Resource 上報自己的狀態，並處理APP部署的任務；
2）Resource：Resource 從 Node 接收心跳，並預聚合心跳後廣播到其他 Resource 節點。Resource 整合所有 Node 組成線上列表，可像無狀態服務一樣水平擴容。為提供業務間隔離性和降低廣播的扇出比，Resource叢集數也會擴充套件；
3）App：App 是執行在 Starlink 上的應用，每個 App 都擁有獨立的資源排程器，這些排程器都從 Resource 獲取全域性的資源檢視，透過樂觀併發搶佔的方式分配資源；
4）Scheduler：Scheduler 負責應用的負載均衡和容災，Scheduler 會根據不同的節點的效能和狀態動態的調整節點的權重，並透過帶權路由演算法來分配請求。
在微信的後臺服務中，每個微服務都是獨立的模組。而面對超大規模的應用，由於 K8s 自身擴縮容效能的限制，往往需要部署多個模組才能滿足一個AI應用，擴縮容速度受限。與K8s 不同的是，Starlink 使用預建立的 Pod，加快了擴縮容的速度，資源遷移變得非常簡單。基於良好的設計，Starlink可以支援單應用百萬節點，樂觀排程也使得排程速度極快，每分鐘可完成數萬節點的排程。Starlink 還可以跨多個資源平臺排程，支援異構機型，不必為每個應用建立多個模組進行部署，大幅提高了內部的資源利用率和資源的週轉效率。
6、 技術挑戰2：不穩定資源下構建穩定服務
6.1 概述
AstraRay 大量接入低價或免費資源，pod 穩定性較差，日常會出現較高的資源驅逐率和亞健康的情況，直接使用會導致服務失敗率高、延時高。另外，用傳統的排程方法排程 AI 計算任務很容易出現計算傾斜，從而導致整體資源利用率低。我們透過更快的容災排程解決服務失敗率高的問題，透過更優的排程演算法來解決服務延時高和資源利用率低的問題。
▲ 圖 14：Starlink 排程流程
6.2 快速容災排程
▲ 圖 15：kubernetes PreStop Hook 機制
我們透過兩個手段來加速容災排程：
1）在資源平臺實際驅逐 pod 之前，透過 K8s 的 PreStop Hook 機制實現服務程式優雅退出，同時Node將自己標記為離線，並透過心跳上報到 Resource。
2）Resouce 透過預聚合廣播，快速將狀態同步到整個 Resouce 叢集，Scheduler 每隔 3s 透過拉取 Resouce 的線上列表來進行動態權重計算，定期更新路由表。最終可以實現在 4s 內將節點驅逐，從而大幅降低了應用的失敗率。
6.3 動態權重SWRR路由演算法
AI 應用往往具有計算量大，單機 QPS 低的特點。在這種服務場景下，微信後臺常用的一致性雜湊已經無法將請求均勻的分發了。除此之外，低優和免費資源因為經常被線上任務搶佔，節點間效能往往參差不齊。我們選用 SWRR(Smooth Weighted Round-Robin)演算法作為基座，並進行最佳化，首次應用到低 QPS 的任務排程系統中，實現請求分佈的快速調整。
演算法步驟如下。
1）更新節點權重(3s一次)：對於每個節點：節點權重=節點核數或卡數∗log(剩餘利用率)∗(當前利用率/節點當前併發)這個公式構建了一個模型，簡單的描述了請求量預期的分佈，節點權重描述的是當前節點處理新增任務的能力，處理能力越高的節點應該分配到更多的請求。
其中：
1) 節點核數或卡數是代表節點的資源總數，資源總數與處理能力成正比，對於不同的GPU，資源總數即不同卡的效能對比係數；
2) log(剩餘利用率)是節點當前剩餘資源，剩餘資源量與處理能力成正比。其中，log是一個經驗值，在log後，演算法在高負載時表現較好；
3) (當前利用率/節點當前併發)本質上是機器效能的體現，假設大盤下每個任務同一時刻的消耗是接近的時，這個公式成立。
2）選擇節點流程：這裡是SWRR的標準流程，因為SWRR演算法的複雜度是O(n)，我們的實現會對效能做一定的最佳化，比如分block，多演算法例項等。
1) 對於每個節點：節點路由權重 = 節點路由權重 + 節點權重；
2) 選擇當前路由權重最大的節點；3) 被選擇的節點的路由權重減去所有節點權重之和。演算法流程樣例，假設{A,B,C}節點權重為{5,1,1}。
最終，我們使用自適應權重的 SWRR 演算法，動態平衡請求分佈，拉平利用率的同時，還降低了請求耗時。
7、 技術挑戰3：降低應用部署的複雜度
▲ 圖 20：AI應用的部署複雜度AI 應用的部署涉及三個方面：多模型擴充套件、多卡型擴充套件、多模組擴充套件（單模組超過 K8s 部署上限），一個超級應用的部署複雜度為 O(n^3)。AstraRay 的創新方案使得一個應用可實現三個維度的擴充套件，將複雜度降低為O(1)，極大提升了 AI 應用部署的效率。
7.1 多模型擴充套件挑戰
多模型擴充套件問題的本質是模型執行環境的動態切換，這裡包含兩個問題：
1）執行時動態切換；
2）模型的快速下發。
動態切換執行時：
▲ 圖 21：Ray動態執行時
我們首先解決執行環境的問題。Ray自身提供RuntimeEnv作為執行環境管理，但Ray的RuntimeEnv無法切換Python版本，且Ray對於Python執行環境之外的依賴，只能依靠機器本身Docker環境，不夠靈活。我們支援了Conda作為Python執行環境的隔離和打包，與Ray本身的Conda不同在於：Ray的Conda要先拉起Ray，而 Ray 的worker節點要求和Ray的頭節點使用相同的版本，導致應用無法切換Python版本。而我們透過在啟動Ray之前初始化執行環境，使每個應用自定義不同的Python版本。具體的操作為：在應用的程式碼打包上傳之前，我們會根據使用者填寫的 requirement.txt，使用conda-pack打包對應的Conda環境，在啟動Ray之前，分發到對應的節點上。其中提前打包可以避免大規模快速擴容對軟體源帶來下載壓力。我們也支援使用者自定義打包例如 TensorRT 等環境，提供更強大的環境自定義能力。
▲ 圖 22：AstraRay 執行時
快速的模型下發：隨著大模型時代的到來，模型檔案變得越來越大，LLM模型有數十GB，下載一個模型需要數十分鐘。Ray可以指定working_dir來分發程式碼和模型，但是Ray單點依賴gcs節點，預設的大小限制也僅僅500MB，無法用於真正的生產環境。為此，我們在Node上嵌入了P2P網路。P2P分為Server端和SDK接入端，server端透過心跳管理P2P節點，並維護叢集中的種子資訊。P2P節點則提供檔案分片的快取和傳輸能力。
▲ 圖 23：P2P server 端架構
▲ 圖 24：P2P sdk 端架構
我們還對P2P的網路和效能做了極致的最佳化：
1）網路打洞能力：面對複雜的網路環境，P2P支援NAT探測打洞，盡最大努力避免網路不通的情況；
2）節點自動限速能力：P2P作為一個嵌入式的元件，要避免節點的頻寬和CPU被P2P程序消耗完，所以節點加入P2P網路時，會對節點進行測速，並設定合適的閾值，避免影響正常服務；
3）全侷限速：即使已經限制了單節點的速度，仍然有可能會因為上層交換機或核心網路頻寬限制，影響到其他服務，支援從服務端下發全網限速，避免影響其他服務；
4）冷啟動和熱點下載加速：一個新的檔案下發時，因為全網都不存在這個檔案，如果按序下載，可能會導致下載緩慢，請求的節點分片集中。透過打亂分片下載的順序，可以將請求分佈到不同的節點。
▲ 圖 25：P2P下載加速
7.2 多模組擴充套件挑戰
▲ 圖 26：Ray 聯邦叢集架構
為了提升 Ray 應用的擴充套件能力，我們透過starlink實現了Ray聯邦叢集架構，每個Ray應用可以擁有多個Ray叢集，單個Ray叢集都擁有完整的功能。使用者可以調整單個Ray叢集的大小，在單個Ray叢集內進行Actor的資源分配，提升應用處理能力，提升資源利用率，實現垂直擴充套件能力；可以透過擴容Ray 叢集數量，實現水平擴充套件。我們還在 Ray 聯邦叢集架構基礎上，增強了 Ray叢集的容災能力，具體策略為：當head node下線，則水平重新擴容一個Ray叢集。當worker node下線，則在這個Ray叢集重新拉起一個worker。透過上述策略，我們使用不穩定的低優資源的情況下，Ray自身架構引起的失敗影響可以降低到最低。
7.3 多卡型擴充套件
▲ 圖 27：TFCC推理執行時
多卡型擴充套件的模型推理部署有三個比較大的挑戰：
1）不同的推理業務形態多樣：引擎種類多模型型別多（pytorch/onnx/tensorrt...）；
2）異構卡型的適配工作繁瑣且重複度高（英偉達/紫霄/華為）；
3）多種引擎支援、模型切換引擎成本高。
我們基於TFCC框架提供標準服務框架，統一了接入模式，透明化了引擎實現，演算法僅需宣告模型，不再需要手寫推理程式碼，同時內化異構卡型適配工作，遮蔽硬體細節，在應用層實現一份程式碼、多處推理，支援靈活多樣的AI應用場景。
8、本文小結
AI 時代的來臨對微信後臺的基礎設施帶來了許多挑戰。我們引入業界先進的Ray作為基座，適配了微信的基礎環境，提供了方便快捷的AI應用開發正規化。同時，在Ray的基礎上，簡化了Ray本身叢集管理的難度，並使用低成本的閒置資源節省了大量的機器成本。AstraRay作為一個剛誕生一年的專案，為微信的AI應用的工程化提供了堅實基礎，並且在持續不斷的最佳化，為將來更多AI應用在微信落地做好了準備。
9、參考資料
[1] Ray on Kubernetes
[2] OpenAI 背書的計算引擎迎里程碑：螞蟻集團成功部署百萬核心計算平臺
[3] 使用 KubeRay 和 Kueue 在 Kubernetes 中託管 Ray 工作負載
[4] Four Reasons Why Leading Companies Are Betting On Ray
[5] The evolution of cluster scheduler architectures
[6] API7 Cloud Integrates with Kubernetes Service Discovery
[7] Upstream: smooth weighted round-robin balancing
[8] Handling files and packages on your cluster with Ray runtime environments
10、微信團隊其它技術文章
《微信技術分享：微信的海量IM聊天訊息序列號生成實踐（演算法原理篇）》
《騰訊技術分享：GIF動圖技術詳解及手機QQ動態表情壓縮技術實踐》
《微信團隊分享：Kotlin漸被認可，Android版微信的技術嚐鮮之旅》
《微信團隊首次揭秘微信紅包演算法，為何你搶到的是0.01元》
《微信團隊分享：極致最佳化，iOS版微信編譯速度3倍提升的實踐總結》
《IM“掃一掃”功能很好做？看看微信“掃一掃識物”的完整技術實現》
《微信團隊分享：微信支付程式碼重構帶來的移動端軟體架構上的思考》
《IM開發寶典：史上最全，微信各種功能引數和邏輯規則資料彙總》
《微信團隊分享：微信直播聊天室單房間1500萬線上的訊息架構演進之路》
《企業微信的IM架構設計揭秘：訊息模型、萬人群、已讀回執、訊息撤回等》
《IM全文檢索技術專題(四)：微信iOS端的最新全文檢索技術最佳化實踐》
《微信團隊分享：微信後臺在海量併發請求下是如何做到不崩潰的》
《微信Windows端IM訊息資料庫的最佳化實踐：查詢慢、體積大、檔案損壞等》
《微信技術分享：揭秘微信後臺安全特徵資料倉儲的架構設計》
《企業微信針對百萬級組織架構的客戶端效能最佳化實踐》
《揭秘企業微信是如何支援超大規模IM組織架構的——技術解讀四維關係鏈》
《微信團隊分享：詳解iOS版微信影片號直播中因幀率異常導致的功耗問題》
《微信團隊分享：微信後端海量資料查詢從1000ms降到100ms的技術實踐》
《大型IM工程重構實踐：企業微信Android端的重構之路》
《IM技術乾貨：假如你來設計微信的群聊，你該怎麼設計？》
《微信團隊分享：來看看微信十年前的IM訊息收發架構，你做到了嗎》
（本文已同步釋出於：http://www.52im.net/thread-4731-1-1.html）