微信後臺基於 Ray 的分佈式 AI 計算技術實踐

一、引言

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

二、本文目錄

_1）_引言

_2）_爲何在 AI 計算中引入 Ray？

**3）**微信基於 Ray 的 AstraRay 平臺

_4）_AstraRay 平臺架構概覽

_5）_技術挑戰 1：單集羣支持百萬級計算節點

• 5.1 架構選擇

• 5.2 Starlink 調度

6） 技術挑戰 2：不穩定資源下構建穩定服務

• 6.1 概述

• 6.2 快速容災調度

• 6.3 動態權重 SWRR 路由算法

7） 技術挑戰 3：降低應用部署的複雜度

• 7.1 多模型擴展挑戰

• 7.2 多模塊擴展挑戰

• 7.3 多卡型擴展

_8）_本文小結

_9）_參考資料

_10）_微信團隊其它技術文章

三、爲何在 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 計算平臺的分佈式底座。

四、微信基於 Ray 的 AstraRay 平臺

P6n 是基於 Kubernetes 微服務部署平臺，通過自動化編排和彈性擴縮容機制，很好的解決了在線高實時的後臺服務運維自動化問題，但不支持大規模的批處理服務，單應用多模型的部署複雜度較高，機器成本較高，不適合 “在離線一體” 的 AI 計算場景。
Gemini 是基於 kubernetes 的大數據平臺，適合處理離線大規模的數據清洗和模型訓練，但是由於調度的實時性不夠，不適合高實時性、高吞吐的和高可靠的 AI 計算場景。
Astra 平臺要實現高實時、高吞吐、高可靠、低成本的 AI 計算平臺，需要解決如下幾個核心問題。
比如：

• _1）_爲了低成本，需要支持各種異構資源擴展；

• _2）_爲了高吞吐，支持超大規模資源調度；

• _3）_降低單應用多模型的部署複雜度。

我們基於 Ray 計算底座，解決了上述三個核心問題，構建出適合 AI 計算平臺：AstraRay，在微信內進行了大規模 AI 應用部署的實踐。
AstraRay 相比社區版本 Ray(KubeRay) 有以下改進：

五、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 集羣管理的問題，通過負載均衡和容災調度解決了不穩定資源構建穩定服務的挑戰，同時通過應用調度解決了多模型應用部署複雜度的問題。接下來詳細介紹我們如何應對這三個技術挑戰。

六、全文鏈接

即時通訊網 (52im.net) 社區鏈接：http://www.52im.net/thread-4731-1-1.html，或點擊下文的 “閱讀原文”！

七、參考文獻

**[**1] Ray on Kubernetes

https://docs.ray.io/en/latest/cluster/kubernetes/index.html

[2] OpenAI 背書的計算引擎迎里程碑：螞蟻集團成功部署百萬核心計算平臺

https://www.infoq.cn/article/vt4lewlrgumufibrulhz

[3] 使用 KubeRay 和 Kueue 在 Kubernetes 中託管 Ray 工作負載

https://juejin.cn/post/7313601254365691941

[4] Four Reasons Why Leading Companies Are Betting On Ray

https://www.anyscale.com/blog/four-reasons-why-leading-companies-are-betting-on-ray

[5] The evolution of cluster scheduler architectures

https://www.cl.cam.ac.uk/research/srg/netos/camsas/blog/2016-03-09-scheduler-architectures.html

[6] API7 Cloud Integrates with Kubernetes Service Discovery

https://api7.ai/blog/api7-cloud-integrates-kubernetes-service-discovery

[7] Upstream: smooth weighted round-robin balancing

https://github.com/nginx/nginx/commit/52327e0627f49dbda1e8db695e63a4b0af4448b1

[8] Handling files and packages on your cluster with Ray runtime environments

https://www.anyscale.com/blog/handling-files-and-packages-on-your-cluster-with-ray-runtime-environments

本文由 Readfog 進行 AMP 轉碼，版權歸原作者所有。
來源：https://mp.weixin.qq.com/s/7WbryLDoVb7F2uyqmsx63A