相信很多人都了解 Ray，它是 OpenAI 用于大模型訓(xùn)練的底層分布式框架。螞蟻集團(tuán)早在很久之前就加入了 Ray，且今年 Ray 的 CEO 在夏季發(fā)布會上也提到，螞蟻是第一個(gè)正式使用并協(xié)作開發(fā) Ray 的團(tuán)隊(duì)。我們貢獻(xiàn)了超過 26% 的 Ray 核心代碼，是全球第二大貢獻(xiàn)團(tuán)隊(duì)。目前，螞蟻集團(tuán)在線上運(yùn)營著超過 150 萬 CPU 核心，規(guī)模已相當(dāng)龐大，同時(shí)我們也在運(yùn)營 Ray 在中國的社區(qū)。

簡單介紹一下 Ray 在螞蟻內(nèi)部的發(fā)展情況。我們在 2017 年成立了 Ray 團(tuán)隊(duì)，并于 2018 年推出了首個(gè)業(yè)務(wù)場景流圖計(jì)算引擎 Geaflow。在 2018 年至 2022 年間的大數(shù)據(jù)時(shí)代，我們基于 Ray 開發(fā)了多個(gè)計(jì)算引擎，如用于流計(jì)算和機(jī)器學(xué)習(xí)訓(xùn)練的 Realtime、Mobius 開源引擎，以及在線推理和科學(xué)計(jì)算引擎 Mars。同時(shí)，我們也貢獻(xiàn)了 Multi-Tenant 架構(gòu)，Ray 社區(qū)今年才開始考慮這種架構(gòu)，而螞蟻因線上集群規(guī)模龐大，早已開始考慮多租戶。

在 2023 到 2024 年大模型時(shí)代，我們在美國完成了一項(xiàng)工作 Unified AI Serving，將離線、在線與 AI 推理、AI 部署整合為一個(gè)框架，這是我們 150 萬內(nèi)核業(yè)務(wù)的核心場景之一。接下來將介紹我們的最新工作，基于 Ray 構(gòu)建的 AI Agent 框架，主要分為三個(gè)部分：背景、動(dòng)因，以及設(shè)計(jì)與實(shí)現(xiàn)。

一、Background

首先來看一下基于大語言模型的 Agent（LLM-based Agent）通常需要哪些模塊。

• 第一個(gè)是 Profile 模塊，它定義了 Agent 的個(gè)性，即扮演怎樣的角色，比如它可以是一個(gè)溫和的旅游助手，執(zhí)行旅游管理、數(shù)據(jù)分析等任務(wù)。
• 第二個(gè)是 Memory 模塊，包括兩部分：一是 Knowledge，包含行業(yè)知識和先驗(yàn)知識；二是 Experience，記錄 Agent 過去的對話、用戶問題、思考過程及行動(dòng)結(jié)果，這些經(jīng)驗(yàn)將幫助 Agent 改進(jìn)后續(xù)的行為，避免重復(fù)錯(cuò)誤。
• 第三個(gè)是 Planning 模塊，用于將復(fù)雜任務(wù)拆解成更容易執(zhí)行的子任務(wù)。通過這種方式，Agent 可以在文本交互的基礎(chǔ)上完成更廣泛的任務(wù)。Planning 常見的算法有 Chain of Thought和Tree of Thought，這些就像程序設(shè)計(jì)中的流程圖，用來拆解復(fù)雜問題。
• 第四個(gè)是 Action 模塊，根據(jù)經(jīng)驗(yàn)和規(guī)劃執(zhí)行實(shí)際任務(wù)。與大模型不同，Agent 不僅僅是文本或圖像輸入輸出，而是能對現(xiàn)實(shí)世界產(chǎn)生實(shí)際影響。Action 模塊的一個(gè)關(guān)鍵功能是 Function Calling，讓模型調(diào)用外部功能，甚至在某些場景中與機(jī)械臂等物理設(shè)備交互。

這四個(gè)模塊是我們認(rèn)為一個(gè)基于大語言模型的 Agent 所需的核心組件。

我們來看一個(gè)簡單的基于 Agent 實(shí)現(xiàn)的 RAG 流程。這個(gè) RAG 不同于傳統(tǒng)的計(jì)算圖或工作流，而是通過 Agent 來實(shí)現(xiàn)的。

• 首先，Agent 會從用戶那里獲取任務(wù)，例如用戶要求從文檔中學(xué)習(xí)知識，提供了文檔鏈接。
• 第二步，模型進(jìn)入思考階段，決定如何開展任務(wù)。這里可以使用 planning 模塊中的算法，如 React 或 Chain of Thought。在這個(gè)流程中，我們使用了 React 策略，即思考一步，執(zhí)行一步。
• 第三步，模型決定采取具體行動(dòng)。對于 RAG 場景，可以使用工具如 LlamaIndex 或 LangChain 對文檔進(jìn)行解析，或調(diào)用 Lucene 和 VectorDB 進(jìn)行語義檢索，甚至進(jìn)行實(shí)時(shí)搜索。
• 第四步，將初始思考、選擇的行動(dòng)及其結(jié)果作為三元組（triplet）存儲到 Memory 中，來指導(dǎo)后續(xù)工作。
• 第五步，繼續(xù)循環(huán)思考和執(zhí)行，直到完成用戶任務(wù)。

特別需要注意的是，右邊的工具分為三類：紅色的 LlamaIndex 用于計(jì)算（CPU/GPU 密集型），綠色的 DB 和 Lucene 訪問為 Disk I/O，藍(lán)色的則是網(wǎng)絡(luò)訪問。Agent 任務(wù)中涉及這三種不同的計(jì)算任務(wù)，傳統(tǒng)模式可能只處理 CPU 或 CPU+IO，但在這里，一個(gè)單任務(wù)可能包含復(fù)雜的混合工作負(fù)載。

在完成 naive RAG Agent 的實(shí)現(xiàn)后，接下來就是部署上線。首先，我們將程序封裝成一個(gè)服務(wù)，放入 Docker 中，然后進(jìn)行部署。我們的 Agent 平臺已經(jīng)有大約 500 個(gè) Agent，開發(fā)者可以在平臺上使用畫布構(gòu)建并發(fā)布他們的 Agent。上圖中就是我們的一些案例。

最初，我們的做法很簡單，但很快收到了很多用戶抱怨。用戶常見的問題包括：不清楚為什么 APP 或 Pod 崩潰、缺乏監(jiān)控 metrics、沒有工作負(fù)載監(jiān)控、無法控制流量等。此外，由于混合工作負(fù)載，GPU 利用率通常很低?？偨Y(jié)起來，這種簡單的做法顯然不適合生產(chǎn)環(huán)境。因此，我們決定將 Agent 平臺分布式化，進(jìn)行生產(chǎn)化改造。在這個(gè)過程中，我們發(fā)現(xiàn)場景中確實(shí)有許多獨(dú)特的挑戰(zhàn)。

二、Motivation

相較于傳統(tǒng)計(jì)算或服務(wù)型 APP，我們認(rèn)為 Agent 應(yīng)用具有高度創(chuàng)造性，不斷涌現(xiàn)出革命性的創(chuàng)意，這也是近年來大模型和 Agent 概念火熱的原因。大家發(fā)現(xiàn)這些應(yīng)用非常有趣，并且持續(xù)有新創(chuàng)意涌現(xiàn)，每個(gè)創(chuàng)意都需要快速低成本驗(yàn)證，這在創(chuàng)業(yè)和商業(yè)中尤為重要。

一旦 PoC 驗(yàn)證有效，就希望能迅速部署到線上，但這一步非常困難。上線應(yīng)用涉及眾多模塊：服務(wù)、數(shù)據(jù)庫、數(shù)據(jù)來源、分布式框架等，對大公司而言，需要與多個(gè)團(tuán)隊(duì)合作，小公司雖不需多個(gè)團(tuán)隊(duì)，但仍需掌握所有技術(shù)。因此，PoC 到上線的過程繁瑣且漫長。此外，用戶任務(wù)可能同時(shí)涉及 GPU、CPU 和 Service Calling，且 Agent 覆蓋不同業(yè)務(wù)場景，每個(gè)場景的技術(shù)棧不同，需要大量兼容適配工作。

在分布式 Ray 生產(chǎn)化 Agent 過程中，我們對比了三個(gè)技術(shù)棧層級：底層的 Kubernetes（簡稱 K8S）、上層的算法庫 LangChain 和中間層的 Ray。Ray 之所以居于中間層，是因?yàn)樗w豐富的 AI 生態(tài)，如 Ray Data 用于數(shù)據(jù)處理和服務(wù)化，以及強(qiáng)化學(xué)習(xí)和訓(xùn)練。Ray 緊貼 K8S，為分布式執(zhí)行調(diào)度提供支持。

Kubernetes 的優(yōu)勢在于底層 API 的靈活性，允許編寫各種 CRD 并整合多種硬件，提供資源控制的完整性；但其缺點(diǎn)是從零開發(fā) AI 程序非常復(fù)雜。我們與基于 K8S 的團(tuán)隊(duì)交流發(fā)現(xiàn)，許多機(jī)器學(xué)習(xí)工程師并不熟悉 K8S。在螞蟻集團(tuán)，算法工程師離掌握 K8S 還有較遠(yuǎn)距離，這種學(xué)習(xí)曲線使得我們不太可能直接讓終端用戶使用 K8S。此外，Kubernetes 的 AI 生態(tài)相對簡陋，主要因其重心仍在底層。

LangChain 作為上層算法庫，優(yōu)點(diǎn)在于豐富且易用，類似的 LlamaIndex 庫也非常易上手。它甚至提供 UI 功能，可以一站式開發(fā) Agent 或大模型。然而，其缺點(diǎn)是僅為單機(jī) API。在我們上線簡易版 LangChain 后，發(fā)現(xiàn)面臨許多生產(chǎn)化問題，每個(gè)問題都需解決。

接下來我們考慮使用 Ray。Ray 的優(yōu)勢在于提供了一站式的工具箱，支持 AI 工作負(fù)載的數(shù)據(jù)處理、訓(xùn)練和推理。Ray 的一個(gè)主要功能是輕松將本地代碼轉(zhuǎn)為分布式代碼，只需簡單加注解即可實(shí)現(xiàn)遠(yuǎn)程進(jìn)程。并且它在異構(gòu)資源調(diào)度方面表現(xiàn)出色，尤其在 CPU 與 GPU 混合的數(shù)據(jù)處理中性能優(yōu)于 Spark。Ray 不綁定特定計(jì)算范式（如 MapReduce 或圖計(jì)算），而是采用純分布式面向?qū)ο缶幊?，提供了高度靈活性。

不過，Ray 也有一些不足：在資源管控上不如 K8S 靈活，因?yàn)樗窃?K8S 之上再加一層；另外，目前 Ray 在大模型上沒有特定 API，一些外部組件如文件傳輸，需要在 Ray API 中額外封裝。

綜上，我們最終選擇了 Ray，同時(shí)也考慮了一些額外的加分項(xiàng)。例如，Ray 的 RuntimeEnv 功能提供了運(yùn)行時(shí)沙箱，利用 container 執(zhí)行用戶代碼，非常適合大模型場景中的代碼解釋器。它能夠直接啟動(dòng)一個(gè) Docker 容器，而在其他團(tuán)隊(duì)中需要額外工作。另一個(gè)優(yōu)勢是 Ray 的面向?qū)ο缶幊棠Ｊ?，與 Agent 的工作模式相似。Agent 的個(gè)性和配置可以視為對象的靜態(tài)資源，操作為成員函數(shù)，記憶則是運(yùn)行時(shí)狀態(tài)，因此 Agent 很像一個(gè)對象。

決定使用 Ray 后，我們自然開始開發(fā)一個(gè)基于 Ray 的 Agent 框架。主要考慮點(diǎn)如下：①該框架需提供 Agent 的 API；②利用 Ray 實(shí)現(xiàn)從本地代碼到支持異構(gòu)資源的分布式代碼的擴(kuò)展；③在多 Agent 場景中，每個(gè) Agent 都是一個(gè)分布式進(jìn)程，我們需要一個(gè)框架來協(xié)調(diào)這些進(jìn)程，即所謂的 environment；④要兼容不同的庫，如 MetaGPT 和 AutoGen；⑤希望利用 Ray 的沙箱（sandbox）、批處理能力和跨源調(diào)度功能。

三、Design & Impl.

在軟件技術(shù)棧中，我們使用 Ray 分為以下幾層：最上層是業(yè)務(wù)層，包括開發(fā)者已編寫的 Agent Apps。其下是 Agent Crafting Platform，這是開發(fā)者構(gòu)建 Agents 的平臺。再下一層是算法庫，如 LangChain 和 MetaGPT，提供與大模型相關(guān)的算法和文檔解析功能，這些不在我們的框架實(shí)現(xiàn)范圍內(nèi)。

在這三個(gè)業(yè)務(wù)層之下是 Ragent，它提供分布式 Agent SDK 和執(zhí)行層，支持 Agent 所需的工具、記憶、環(huán)境管理、分布式通信及部署等功能。Ragent 不實(shí)現(xiàn)具體算法，而是將用戶代碼分布式化，依賴 Ray 的核心概念如 Task、Actor 和 Object，利用 Ray 的編程原語實(shí)現(xiàn)分布式面向?qū)ο缶幊獭ay Data 用于批處理，例如統(tǒng)一清洗和解析用戶文檔，類似于 Spark。Agent 的服務(wù)化通過 reserve 實(shí)現(xiàn)。使用 Ray 后，我們獲得高編程性能，并具備 Failover（故障轉(zhuǎn)移）能力、分級調(diào)度和共享內(nèi)存的優(yōu)勢。最底層是 K8S，用于資源管理和調(diào)度。

使用 Ragent 編寫一個(gè) Agent 時(shí)，首先要了解 Ragent 提供的 Agent 概念，這是一個(gè)具備 Failover 能力的基礎(chǔ)單元，并內(nèi)置消息隊(duì)列和 Memory。Memory 是 Agent 內(nèi)部的記憶。由于存在多種 Planning 策略，Ragent 內(nèi)置了 ReAct 算法，用于重復(fù)的 think 和 act 過程。如上圖的簡化版代碼中，通過一個(gè) while True 循環(huán)，每次先進(jìn)行 reasoning（think），獲取 thought 和預(yù)期執(zhí)行的 action，然后實(shí)際調(diào)用 action，持續(xù)在循環(huán)中執(zhí)行。

Agent 的能力主要來自于其具備的工具，如文檔清洗或制定旅游攻略的能力，這使得 Agent 能夠在功能上超越單純的大模型，因此 Tool 部分至關(guān)重要。要為 Agent 定制、加載或注冊一個(gè) tool，我們可以通過注解實(shí)現(xiàn)。在 Python 代碼中，比如對 index_doc 函數(shù)輸入文檔，使用 Ray Data 和 LangChain 進(jìn)行處理。用 tool 注解這個(gè)函數(shù)即可將其注冊到 Agent 中，使 Agent 理解。注冊過程中，需要用自然語言描述其功能、輸入輸出及用途，類似于編程中的 doc string。注冊的每個(gè) tool 會作為 prompt 的一部分輸入到大模型，使得整個(gè) Agent 相當(dāng)于一個(gè) Ray Actor。

實(shí)現(xiàn)一個(gè)簡易版的 RAG Agent 在我們的框架中非常簡單，只需幾行代碼。首先，引入必要的庫后，在 main 函數(shù)中初始化一個(gè) Agent，使用內(nèi)置的 ReAct 算法，并指定大模型為 Qwen。在 profile 中賦予靜態(tài)資源后，為 Agent 注冊多個(gè)工具，如 Apache Lucene 和 LlamaIndex。我們注冊的 action 包括使用 LlamaIndex 進(jìn)行文檔索引，Lucene 用于 Elasticsearch，以及 DB 進(jìn)行語義搜索，這些工具便可實(shí)現(xiàn)一個(gè)簡單的 RAG。

完成初始化和注冊后，Agent 便可與用戶交互。用戶的第一個(gè) involve 是從 Ray 文檔中學(xué)習(xí)最新功能，此時(shí) Agent 會調(diào)用工具，將用戶提供的文檔索引到 DB 中作為知識的一部分。工具使用 Ray Data 實(shí)現(xiàn)，因此在執(zhí)行階段，Agent 將其作為 Ray Data 的作業(yè)提交，批處理輸出到 DB 中每個(gè)文檔的知識。

接下來，用戶可對文檔提問，如詢問 Ray 最新版本的 Ray Data 功能，此時(shí) Agent 會從 DB 中進(jìn)行字符和語義檢索，經(jīng)過大模型處理后返回結(jié)果。整個(gè)過程只需幾行代碼即可完成一個(gè) RAG Agent，不過這是簡易版的。

接下來，我們看一個(gè) Multi-agent 的例子。左邊的圖是 MetaGPT 的實(shí)現(xiàn)，展示了如何利用多個(gè) Agent 構(gòu)建一個(gè)軟件公司。每個(gè) Agent 承擔(dān)不同職責(zé)，如產(chǎn)品經(jīng)理、架構(gòu)師、代碼工程師和測試工程師。對于任務(wù)如編寫貪吃蛇程序，Agent 按順序協(xié)作。產(chǎn)品經(jīng)理用工具生成設(shè)計(jì)圖給架構(gòu)師，架構(gòu)師再創(chuàng)建技術(shù)架構(gòu)圖。工程師根據(jù)用戶需求編寫接口和實(shí)現(xiàn)，再交給測試進(jìn)行 UT，直到程序完成。左圖展示了流程，右圖是用戶交互，比如編寫 FlappyBird，此例子被 Ragent 框架分布式化。

在該框架中，我們實(shí)現(xiàn)了一個(gè) environment 組件，用于 task 追蹤，包含用戶任務(wù)和每個(gè) Agent 的任務(wù)。它構(gòu)建 workflow，通過 message queue 與 Agent 通信，并保存對話歷史。每個(gè) Agent 作為遠(yuǎn)程進(jìn)程，由 Agent handler 管理。在代碼中，我們先初始化 environment（紫色為  MetaGPT 代碼，橙色為框架代碼），然后初始化架構(gòu)師、產(chǎn)品經(jīng)理和 coder 等 agent，為 MetaGPT 代碼進(jìn)行適配。

初始化四個(gè) Agent 后，將其加入 environment。每個(gè) Agent 描述功能已在 profile 和 system prompt 中定義，注冊后 environment 知道各 Agent 職責(zé)。目前還需手動(dòng)指定 Agent 交互順序，environment 尚不能自動(dòng)選擇。加入后，環(huán)境運(yùn)行應(yīng)用，如編寫 FlappyBird 或貪吃蛇。在實(shí)踐中，GPT-4 效果較好，能構(gòu)建設(shè)計(jì)圖和部分可運(yùn)行代碼，但其他模型仍難以實(shí)現(xiàn)復(fù)雜應(yīng)用，通常在第一輪生成輸出。以上是 Ragent 框架在 Multi-agent 場景的應(yīng)用。

接下來是我們未來的一些工作。首先是 Agent Mesh。目前，Agent 框架眾多，但缺乏統(tǒng)一的通信和流程標(biāo)準(zhǔn)。我們希望通過 Agent Protocol 項(xiàng)目，制定一個(gè)協(xié)議，整合不同框架，實(shí)現(xiàn)類似服務(wù)網(wǎng)格的通信環(huán)境。右圖展示的是我們正在開發(fā)的離在線一體架構(gòu)，這在非 Agent 場景下已實(shí)現(xiàn)。由于底層執(zhí)行層都是 Ray，無論在線還是離線，技術(shù)棧相同，我們無需特別定制，能夠結(jié)合使用。

在 Agent 場景中，文檔處理等純離線操作通過 Ray Data 實(shí)現(xiàn)，第一步可用 Ray Data pipeline 完成離線工作。對于單 Agent 或多 Agent 的二三步，可以實(shí)現(xiàn)服務(wù)化，每個(gè)進(jìn)程用 Agent Protocol 封裝，實(shí)現(xiàn)互通信。這在 Agent 場景中尚未完成，但在計(jì)劃中。

我們還需關(guān)注底層硬件。目前不需要大量 GPU，但有廠商和開源社區(qū)希望支持更多 GPU，如 NPU。以上即是我們的工作計(jì)劃。