作者：微信Astra團隊
一、背景微信存在大量AI計算的應用場景，主要分為三種：流量分發(fā)、產(chǎn)品運營和內(nèi)容創(chuàng)作。流量分發(fā)場景中的 AI 計算主要用于搜索、廣告、推薦場景的核心特征生產(chǎn)，產(chǎn)品運營相關的 AI 計算主要用于產(chǎn)品功能相關和內(nèi)容運營相關（低質(zhì)、優(yōu)質(zhì)、生態(tài)建設），由于大模型的興起，AIGC 相關的文生圖、圖生圖、AI 特效等內(nèi)容創(chuàng)作場景的 AI 計算也有了較多的落地。目前AI 計算幾乎覆蓋了微信的所有業(yè)務場景。
圖 1：微信內(nèi) AI 計算應用場景然而，我們在使用微信已有的后臺基礎設施實現(xiàn)AI應用時遇到各種問題。在資源層面，AI應用屬于計算密集型，計算復雜度高，需要大量資源，直接使用在線資源會導致成本過高；在部署層面，微信后臺常見的部署平臺更適合部署I/O密集、高并發(fā)、高請求量的微服務，而AI應用則需要適配大量異構硬件和異構資源平臺，部署復雜度呈指數(shù)級上升；在應用編排層面，直接通過消息隊列等基礎組件解決復雜特征依賴及相關異步過程，開發(fā)效率低，變更風險高，可觀測性差；在平臺層面，由于缺乏平臺支撐，算法迭代速度慢，模型能力使用門檻高。因此，微信亟需一個低成本、高效率、低門檻的AI計算平臺來解決上述問題。
圖 2：微信內(nèi)原有基礎設施比如，OCR作為視頻號推薦和視頻號搜索依賴的一個重要特征，計算量非常大，需要超過100 萬核的CPU計算資源，同時對實時性和可靠性的要求很高，需要在 1 分鐘內(nèi)完成特征生成。P6n平臺適合做高實時(毫秒級響應)的在線任務，實時性上可以滿足需求，但固定部署的資源成本較高，多模型部署復雜度高，不符合需求。Gemini 平臺更適合做大規(guī)模長時間的離線任務，在實時性和可靠性上不滿足需求。我們需要一個高實時（10 秒級響應），支持大規(guī)模異構資源部署，低成本和高可靠的近線任務平臺。
二、為何在AI計算中引入Ray？
圖 3：使用 Ray 構建 AI 計算的企業(yè)[1]Ray是一個通用的分布式計算引擎，2016年開源于加州大學伯克利分校 RISELab，是發(fā)展最快的計算引擎之一。目前已經(jīng)廣泛應用于OpenAI、螞蟻、字節(jié)和華為等公司，是新一代的明星計算框架。
首先，編寫分布式計算既簡單又直觀。開發(fā)者不必了解所有通信和調(diào)度細節(jié)，也不必對此進行推理。借助 Ray 的簡單原語，可以將任何 Python 函數(shù)或類轉換為分布式執(zhí)行：只需添加一個裝飾器，就大功告成了。Ray 的分布式API 很簡單，所有復雜性都由 Ray 的執(zhí)行框架處理。函數(shù)將被安排為無狀態(tài)任務執(zhí)行，而類將是一個有狀態(tài)的遠程服務。
defdetect(image_data):model=load_detect_model()returnmodel(image_data)defrecognize(det_result):model=load_recognize_model()returnmodel(det_result)defocr(image_data):det_result=detect(image_data)returnrecognize(det_result)image_data=load_image_data()ocr_result=ocr(image_data)以上是一個圖片ocr本地執(zhí)行的 python 腳本，如果使用微服務部署，因為模型過大，單機顯存不夠，無法加載所有模型，則需要部署三個微服務模塊：detect、recognize和ocr，應用部署的復雜度較高。
@ray.remote(num_gpus=1,num_cpus=16)defdetect(image_data):model=load_detect_model()returnmodel(image_data)@ray.remote(num_gpus=2,num_cpus=16)defrecognize(detect_result):model=load_recognize_model()returnmodel(detect_result)@ray.remote(num_cpus=4)defocr(image_data):det_result=detect.remote(image_data)returnrecognize.remote(det_result)image_data=load_image_data()ocr_result=ocr.remote(image_data)如果使用 ray 來做 ocr 推理，只需要添加裝飾器@remote，指定模型使用的 cpu 和 gpu 資源數(shù)，通過一個python 腳本即可完成ocr應用的部署，效率提升至少一個數(shù)量級。
圖 4：Ray AIR 如何以簡單的方式統(tǒng)一 ML 庫[2]其次，大多數(shù)流行的 ML 庫都與 Ray 有很強的集成性，而且 Ray 的原生庫也整合了這些庫。例如，開發(fā)者可以輕松地將 XGBBoost 與 Ray Train 結合使用，可以輕松地將 HuggingFace 與 Ray Serve 結合使用�；蛘�，可以輕松地將 PyTorch 和 TensorFlow 與 Ray Train 結合使用。簡而言之，它擁有豐富的集成生態(tài)系統(tǒng)，不僅與 ML 庫集成，還與其他工具和框架集成。
第三，開發(fā)人員可以使用筆記本電腦進行開發(fā)。當你想將其擴展到 Ray 集群時，只需更改一行代碼或不更改任何代碼即可輕松完成。
RAY_ADDRESS=ray://
三、Ray在微信AI計算中的大規(guī)模實踐
P6n是基于 Kubernetes微服務部署平臺，通過自動化編排和彈性擴縮容機制，很好的解決了在線高實時的后臺服務運維自動化問題，但不支持大規(guī)模的批處理服務，單應用多模型的部署復雜度較高，機器成本較高，不適合“在離線一體”的 AI計算場景。Gemini 是基于 kubernetes 的大數(shù)據(jù)平臺，適合處理離線大規(guī)模的數(shù)據(jù)清洗和模型訓練，但是由于調(diào)度的實時性不夠，不適合高實時性、高吞吐的和高可靠的AI計算場景。Astra 平臺要實現(xiàn)高實時、高吞吐、高可靠、低成本的 AI 計算平臺，需要解決如下幾個核心問題：1.為了低成本，需要支持各種異構資源擴展；2.為了高吞吐，支持超大規(guī)模資源調(diào)度；3.降低單應用多模型的部署復雜度；
我們基于 Ray 計算底座，解決了上述三個核心問題，構建出適合 AI 計算平臺：AstraRay，在微信內(nèi)進行了大規(guī)模 AI 應用部署的實踐。AstraRay 相比社區(qū)版本Ray(KubeRay) 有以下改進：
3.1. AstraRay架構
圖 5：kuberay 架構[3]
圖 6：KubeRay 提交任務流程[4]業(yè)界使用社區(qū)成熟的 KubeRay 方案，通過 Ray 和 K8s 結合，提供了易用、高可用、高伸縮的云原生 Ray 集群服務，可以滿足中小規(guī)模 AI 應用的需求。但它有集群規(guī)模小（最大僅支持數(shù)千個節(jié)點），異構資源擴展困難（單個 ray 集群只能部署在一個 k8s 集群，不支持聯(lián)邦k8s 集群）和伸縮慢（受限于 K8s 的擴縮容速度）的問題，不適合微信內(nèi)超大規(guī)模 AI 應用的需求。
圖 7：AstraRay 整體架構我們在落地 Ray 的過程中遇到了三個核心技術挑戰(zhàn)：1.百萬級 pod 的集群管理：在視頻號業(yè)務場景中，有超過百萬核的超級應用，已經(jīng)遠超 K8s 集群上限，我們希望單個 Ray 應用能支持百萬級別的 pod 的擴展。
2.不穩(wěn)定資源下構建穩(wěn)定服務：由于 AI 計算的資源消耗大，為了降低成本，我們大量使用了低成本、閑置，但穩(wěn)定性差的計算資源。我們希望可以在不穩(wěn)定資源上提供可靠穩(wěn)定的服務。
3.降低應用部署的復雜度：微信內(nèi) AI 應用遇到模型、硬件、模塊三種維度的異構問題，部署復雜度高。我們希望使用統(tǒng)一的應用維度來簡化應用部署，即將 O(n^3) 復雜度降低為 O(1)。
Astra 的部署系統(tǒng)架構如上圖，在 Poseidon/算力/太極/Gemini 等多個資源平臺基礎上擴展多個tke模塊，組成擁有數(shù)百萬核CPU、萬卡GPU級別的超大集群。我們通過服務發(fā)現(xiàn)的架構設計，解決了百萬級pod集群管理的問題，通過負載均衡和容災調(diào)度解決了不穩(wěn)定資源構建穩(wěn)定服務的挑戰(zhàn)，同時通過應用調(diào)度解決了多模型應用部署復雜度的問題。接下來詳細介紹我們?nèi)绾螒獙@三個技術挑戰(zhàn)。
3.2. 單集群支持百萬級計算節(jié)點的挑戰(zhàn)3.2.1 架構選擇
圖 8：集群調(diào)度架構分類[5]業(yè)界系統(tǒng)的調(diào)度架構主要分為四類：單體調(diào)度、兩層調(diào)度、共享調(diào)度和混合調(diào)度。這些調(diào)度架構的本質(zhì)區(qū)別其實只有兩點：1.調(diào)度時調(diào)度器是否擁有全局的資源視圖；2.不同的應用是否擁有多個資源調(diào)度器。
單體調(diào)度顧名思義，即只有一個調(diào)度器，調(diào)度器擁有全局資源視圖的架構，Google Borg 和 K8s 都采用這個架構。單體架構的好處是，所有的任務都由唯一的調(diào)度器處理，調(diào)度器可以充分的考慮全局的資源使用情況，能方便的做出最優(yōu)調(diào)度。但由于調(diào)度架構的限制，集群性能受限于單體的性能，無法支撐過大的集群。
兩層調(diào)度擁有多個調(diào)度器，Apache Mesos 和 Hadoop YARN 都采用這個架構。兩層調(diào)度中，每個應用的調(diào)度器首先向中心節(jié)點獲取資源，再將其分配給應用中的各個任務。兩層調(diào)度解決了單體調(diào)度的性能問題，但是調(diào)度器僅擁有局部資源視圖，無法做出最優(yōu)調(diào)度。
共享調(diào)度擁有多個調(diào)度器，每個調(diào)度器擁有全局資源視圖，Omega 采用了這個架構。共享調(diào)度方案中，每個調(diào)度器都可以并發(fā)地從整個資源池中申請資源，解決了性能問題和最優(yōu)調(diào)度問題，且可以支持較大集群。因此，AstraRay 選擇共享調(diào)度來支持超大規(guī)模的資源管理。調(diào)度器間資源申請沖突可通過悲觀鎖或樂觀鎖來解決，AstraRay 實現(xiàn)了基于樂觀鎖的方案，出現(xiàn)沖突后再處理，無需中心節(jié)點，并發(fā)度更高。
3.2.2 Starlink調(diào)度我們提出了一個新的調(diào)度系統(tǒng) Starlink 來更好適配異構資源和硬件。Starlink采用共享調(diào)度架構，通過樂觀并發(fā)調(diào)度處理沖突，支持部署在任何基礎資源平臺(K8s/Yard/CVM)之上，且允許單個應用運行于多種異構的資源節(jié)點上。
圖 9：Starlink 調(diào)度架構Starlink主要分為四個部分：Node：任意部署了 Starlink 的 Agent 節(jié)點都可以成為 Node，Node 每秒會向Resource 上報自己的狀態(tài)，并處理APP部署的任務。
Resource：Resource 從 Node 接收心跳，并預聚合心跳后廣播到其他 Resource 節(jié)點。Resource 整合所有 Node 組成在線列表，可像無狀態(tài)服務一樣水平擴容。為提供業(yè)務間隔離性和降低廣播的扇出比，Resource集群數(shù)也會擴展。
App：App 是運行在 Starlink 上的應用，每個 App 都擁有獨立的資源調(diào)度器，這些調(diào)度器都從 Resource 獲取全局的資源視圖，通過樂觀并發(fā)搶占的方式分配資源。
Scheduler：Scheduler 負責應用的負載均衡和容災，Scheduler 會根據(jù)不同的節(jié)點的性能和狀態(tài)動態(tài)的調(diào)整節(jié)點的權重，并通過帶權路由算法來分配請求。
在微信的后臺服務中，每個微服務都是獨立的模塊。而面對超大規(guī)模的應用，由于 K8s 自身擴縮容性能的限制，往往需要部署多個模塊才能滿足一個AI應用，擴縮容速度受限。與K8s 不同的是，Starlink 使用預創(chuàng)建的 Pod，加快了擴縮容的速度，資源遷移變得非常簡單�；诹己玫脑O計，Starlink可以支持單應用百萬節(jié)點，樂觀調(diào)度也使得調(diào)度速度極快，每分鐘可完成數(shù)萬節(jié)點的調(diào)度。Starlink 還可以跨多個資源平臺調(diào)度，支持異構機型，不必為每個應用創(chuàng)建多個模塊進行部署，大幅提高了內(nèi)部的資源利用率和資源的周轉效率。
3.3. 不穩(wěn)定資源下構建穩(wěn)定服務的挑戰(zhàn)
AstraRay 大量接入低價或免費資源，pod 穩(wěn)定性較差，日常會出現(xiàn)較高的資源驅(qū)逐率和亞健康的情況，直接使用會導致服務失敗率高、延時高。另外，用傳統(tǒng)的調(diào)度方法調(diào)度 AI 計算任務很容易出現(xiàn)計算傾斜，從而導致整體資源利用率低。我們通過更快的容災調(diào)度解決服務失敗率高的問題，通過更優(yōu)的調(diào)度算法來解決服務延時高和資源利用率低的問題。
圖 10：Starlink 調(diào)度流程3.3.1 快速容災調(diào)度
圖 11：kubernetes PreStop Hook 機制[6]我們通過兩個手段來加速容災調(diào)度：1. 在資源平臺實際驅(qū)逐 pod 之前，通過 K8s 的 PreStop Hook 機制實現(xiàn)服務程序優(yōu)雅退出，同時Node將自己標記為離線，并通過心跳上報到 Resource。2. Resouce 通過預聚合廣播，快速將狀態(tài)同步到整個 Resouce 集群，Scheduler 每隔 3s 通過拉取 Resouce 的在線列表來進行動態(tài)權重計算，定期更新路由表。最終可以實現(xiàn)在 4s 內(nèi)將節(jié)點驅(qū)逐，從而大幅降低了應用的失敗率。
3.3.2 動態(tài)權重SWRR路由算法AI 應用往往具有計算量大，單機 QPS 低的特點。在這種服務場景下，微信后臺常用的一致性哈希已經(jīng)無法將請求均勻的分發(fā)了。除此之外，低優(yōu)和免費資源因為經(jīng)常被在線任務搶占，節(jié)點間性能往往參差不齊。我們選用 SWRR [6](Smooth Weighted Round-Robin)算法作為基座，并進行優(yōu)化，首次應用到低 QPS 的任務調(diào)度系統(tǒng)中，實現(xiàn)請求分布的快速調(diào)整。算法步驟如下：
1.更新節(jié)點權重(3s一次)
對于每個節(jié)點：節(jié)點權重=節(jié)點核數(shù)或卡數(shù)log(剩余利用率)(當前利用率/節(jié)點當前并發(fā))，這個公式構建了一個模型，簡單的描述了請求量預期的分布，節(jié)點權重描述的是當前節(jié)點處理新增任務的能力，處理能力越高的節(jié)點應該分配到更多的請求。其中：1)節(jié)點核數(shù)或卡數(shù)是代表節(jié)點的資源總數(shù)，資源總數(shù)與處理能力成正比，對于不同的GPU，資源總數(shù)即不同卡的性能對比系數(shù)。2)log(剩余利用率)是節(jié)點當前剩余資源，剩余資源量與處理能力成正比。其中，log是一個經(jīng)驗值，在log后，算法在高負載時表現(xiàn)較好。3)(當前利用率/節(jié)點當前并發(fā))本質(zhì)上是機器性能的體現(xiàn)，假設大盤下每個任務同一時刻的消耗是接近的時，這個公式成立。
2.選擇節(jié)點流程
這里是SWRR的標準流程，因為SWRR算法的復雜度是O(n)，我們的實現(xiàn)會對性能做一定的優(yōu)化，比如分block，多算法實例等。1)對于每個節(jié)點：節(jié)點路由權重 = 節(jié)點路由權重 + 節(jié)點權重2)選擇當前路由權重最大的節(jié)點3)被選擇的節(jié)點的路由權重減去所有節(jié)點權重之和算法流程樣例，假設{A,B,C}節(jié)點權重為{5,1,1}
最終，我們使用自適應權重的 SWRR 算法，動態(tài)平衡請求分布，拉平利用率的同時，還降低了請求耗時。
3.4. 降低應用部署的復雜度的挑戰(zhàn)
圖 12：AI應用的部署復雜度AI 應用的部署涉及三個方面：多模型擴展、多卡型擴展、多模塊擴展（單模塊超過 K8s 部署上限），一個超級應用的部署復雜度為 O(n^3)。AstraRay 的創(chuàng)新方案使得一個應用可實現(xiàn)三個維度的擴展，將復雜度降低為O(1)，極大提升了 AI 應用部署的效率。
3.4.1. 多模型擴展挑戰(zhàn)多模型擴展問題的本質(zhì)是模型運行環(huán)境的動態(tài)切換，這里包含兩個問題：
1.運行時動態(tài)切換；2.模型的快速下發(fā)；動態(tài)切換運行時
圖 13：Ray動態(tài)運行時[7]我們首先解決運行環(huán)境的問題。Ray自身提供RuntimeEnv作為運行環(huán)境管理，但Ray的RuntimeEnv無法切換Python版本，且Ray對于Python運行環(huán)境之外的依賴，只能依靠機器本身Docker環(huán)境，不夠靈活。
我們支持了Conda作為Python運行環(huán)境的隔離和打包，與Ray本身的Conda不同在于：Ray的Conda要先拉起Ray，而 Ray 的worker節(jié)點要求和Ray的頭節(jié)點使用相同的版本，導致應用無法切換Python版本。而我們通過在啟動Ray之前初始化運行環(huán)境，使每個應用自定義不同的Pyton版本。
具體的操作為：在應用的代碼打包上傳之前，我們會根據(jù)用戶填寫的 requirement.txt，使用conda-pack打包對應的Conda環(huán)境，在啟動Ray之前，分發(fā)到對應的節(jié)點上。其中提前打包可以避免大規(guī)�？焖贁U容對軟件源帶來下載壓力。我們也支持用戶自定義打包例如 TensorRT 等環(huán)境，提供更強大的環(huán)境自定義能力。
圖 14：AstraRay 運行時快速的模型下發(fā)
隨著大模型時代的到來，模型文件變得越來越大，LLM模型有數(shù)十GB，下載一個模型需要數(shù)十分鐘。Ray可以指定working_dir來分發(fā)代碼和模型，但是Ray單點依賴gcs節(jié)點，默認的大小限制也僅僅500MB，無法用于真正的生產(chǎn)環(huán)境。
為此，我們在Node上嵌入了P2P網(wǎng)絡。P2P分為Server端和SDK接入端，server端通過心跳管理P2P節(jié)點，并維護集群中的種子信息。P2P節(jié)點則提供文件分片的緩存和傳輸能力。
圖 15：P2P server 端架構
圖 16：P2P sdk 端架構我們還對P2P的網(wǎng)絡和性能做了極致的優(yōu)化：1.網(wǎng)絡打洞能力：面對復雜的網(wǎng)絡環(huán)境，P2P支持NAT探測打洞，盡最大努力避免網(wǎng)絡不通的情況。
2.節(jié)點自動限速能力：P2P作為一個嵌入式的組件，要避免節(jié)點的帶寬和CPU被P2P進程消耗完，所以節(jié)點加入P2P網(wǎng)絡時，會對節(jié)點進行測速，并設定合適的閾值，避免影響正常服務。
3.全局限速：即使已經(jīng)限制了單節(jié)點的速度，仍然有可能會因為上層交換機或核心網(wǎng)絡帶寬限制，影響到其他服務，支持從服務端下發(fā)全網(wǎng)限速，避免影響其他服務。
4.冷啟動和熱點下載加速：一個新的文件下發(fā)時，因為全網(wǎng)都不存在這個文件，如果按序下載，可能會導致下載緩慢，請求的節(jié)點分片集中。通過打亂分片下載的順序，可以將請求分布到不同的節(jié)
圖 17：P2P下載加速3.4.2. 多模塊擴展挑戰(zhàn)
圖 18：Ray 聯(lián)邦集群架構為了提升 Ray 應用的擴展能力，我們通過starlink實現(xiàn)了Ray聯(lián)邦集群架構，每個Ray應用可以擁有多個Ray集群，單個Ray集群都擁有完整的功能。用戶可以調(diào)整單個Ray集群的大小，在單個Ray集群內(nèi)進行Actor的資源分配，提升應用處理能力，提升資源利用率，實現(xiàn)垂直擴展能力；可以通過擴容Ray 集群數(shù)量，實現(xiàn)水平擴展。
我們還在 Ray 聯(lián)邦集群架構基礎上，增強了 Ray集群的容災能力，具體策略為：當head node下線，則水平重新擴容一個Ray集群。當worker node下線，則在這個Ray集群重新拉起一個worker。通過上述策略，我們使用不穩(wěn)定的低優(yōu)資源的情況下，Ray自身架構引起的失敗影響可以降低到最低。
3.4.3. 多卡型擴展
圖 19：TFCC推理運行時多卡型擴展的模型推理部署有三個比較大的挑戰(zhàn)：
1.不同的推理業(yè)務形態(tài)多樣：引擎種類多模型類型多（pytorch/onnx/tensorrt...）2.異構卡型的適配工作繁瑣且重復度高（英偉達/紫霄/華為）3.多種引擎支持、模型切換引擎成本高。
我們基于TFCC框架提供標準服務框架，統(tǒng)一了接入模式，透明化了引擎實現(xiàn)，算法僅需聲明模型，不再需要手寫推理代碼，同時內(nèi)化異構卡型適配工作，屏蔽硬件細節(jié)，在應用層實現(xiàn)一份代碼、多處推理，支持靈活多樣的AI應用場景。
四、總結AI 時代的來臨對微信后臺的基礎設施帶來了許多挑戰(zhàn)。我們引入業(yè)界先進的Ray作為基座，適配了微信的基礎環(huán)境，提供了方便快捷的AI應用開發(fā)范式。同時，在Ray的基礎上，簡化了Ray本身集群管理的難度，并使用低成本的閑置資源節(jié)省了大量的機器成本。AstraRay作為一個剛誕生一年的項目，為微信的AI應用的工程化提供了堅實基礎，并且在持續(xù)不斷的優(yōu)化，為將來更多AI應用在微信落地做好了準備。
