微信后團隊分享：微信后臺基于Ray的分布式AI計算技術實踐

Posted on 2024-11-07 11:07 Jack Jiang 閱讀(105) 評論(0) 編輯收藏

本文由微信后臺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 秒級響應），支持大規模異構資源部署，低成本和高可靠的近線任務平臺。

技術交流：

- 移動端IM開發入門文章：《新手入門一篇就夠：從零開發移動端IM》

- 開源IM框架源碼：https://github.com/JackJiang2011/MobileIMSDK（備用地址點此）

（本文已同步發布于：http://www.52im.net/thread-4731-1-1.html）

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 應用的需求。但它有集群規模?。ㄗ畲髢H支持數千個節點），異構資源擴展困難（單個 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，加快了擴縮容的速度，資源遷移變得非常簡單?；诹己玫脑O計，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之前，分發到對應的節點上。其中提前打包可以避免大規?？焖贁U容對軟件源帶來下載壓力。我們也支持用戶自定義打包例如 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應用在微信落地做好了準備。