初窺Ray框架

本文首發於：行者AI

隨着各行各業數字化的不斷推動，AI須要處理的數據愈來愈多，單一服務器已經難以知足當前產業的發展需求，服務器集羣成爲企業用AI處理數據的標配硬件，而分佈式計算成爲人工智能應用的標配軟件。

從圖1能夠看出，現今有不少開源的分佈式計算框架，從模型的訓練、調參到部署；從NLP、CV到RS；這些框架覆蓋到了AI產業生命週期的各個方面。本文就選取其中的Ray框架進行簡單的介紹。

圖1. 各類分佈式計算框架

Ray 是伯克利大學在2017年開源的分佈式計算框架，對應的論文是《Ray: A Distributed Framework for Emerging AI Applications》。強化學習任務須要與環境進行大量的交互（毫秒級），且在時間上支持異構性。該框架專門爲機器學習與強化學習設計，相較於其餘框架，ray具備如下優點：

• 輕量級
• 可快速構建
• 通用性強
• 性能優異

下面就這四個優勢爲你們進行詳細介紹。

1. Ray框架的優點

1.1 輕量級

相較於傳統的分佈式框架（尤爲是hadoop、spark等），Ray能夠直接經過pip進行安裝，且對系統版本無要求。

pip install -U ray

Ray是一個簡單的分佈式策略，而非完整的生態，於是不須要複雜的構建。

另外一方面，輕量而優秀的框架每每能夠做爲企業數據處理的基礎框架，企業不斷在該框架的基礎上增長生態，從而造成企業獨有的應用生態。

1.2 可快速構建

如hadoop等傳統框架，要對原有的單機程序進行分佈式化，須要修改整個代碼邏輯，以MapReduce的編程方案重構各個計算模塊，這使得hadoop等傳統框架有着良好的可編輯性，算法工程師能夠根據業務需求進行詳細的修改。強大的可編輯性也帶來了學習成本高，代碼重構困難等諸多問題。人工智能突飛猛進，模型在不斷更迭，敏捷開發成爲了不少AI企業的開發模式，AI應用的複雜構建會大大影響整個項目的推動。

以下代碼，將一個簡單的單機程序函數，轉換爲Ray分佈式的函數，只是在原有函數的基礎上加入了ray.remote的裝飾器，便完成了分佈式化的工做。

### 原始單機代碼
def f(x):
    return x * x

futures = [f.remote(i) for i in range(4)]
print(ray.get(futures))

### Ray分佈式代碼
import ray
ray.init()
@ray.remote
def f(x):
    return x * x

futures = [f.remote(i) for i in range(4)]
print(ray.get(futures))

1.3 通用性強

近年tensorflow、torch等深度學習框架成爲人工智能應用的模型框架，考慮到產業應用場景，這些框架都給出了各自分佈式訓練和部署的方案，且這些方案的計算資源利用率較高。大型的項目每每由數個算法模型組成，爲了快速開發，算法工程師每每採用開源的代碼構建，而這些開源的代碼採用的深度學習框架極可能互不相同，針對單一框架的分佈式方案難以適用。

除此以外，ONNIX等爲表明的框架，傾向於將全部框架的模型統一到單一的解決方案上，因爲不少前沿的深度學習模型對神經元進行了複雜的修改，沒法適配到通用的算子上，須要算法工程師手寫算子，從而拖慢了開發速度。Ray將機器學習模型、numpy數據計算、單一的函數抽象成通用的計算，實現了對各類深度學習框架、機器學習框架的適配。

另外，Ray對強化學習的應用進行了專門的生態構建。

1.4 性能優異

圖2爲Ray、Horovod以及tensorflow原生的分佈式方案訓練ResNet-101模型的比較，縱軸爲每秒平均迭代的圖片數，能夠看出Ray略微優於Horovod框架。

圖2. 分佈式訓練速度比較

圖3爲Clipper和Ray在模型調用上吞吐量的比較，二者均用同一網絡模型，能夠看出Ray優於Clipper。

圖3. 分佈式部署吞吐量比較

Ray並無作到每一個分佈式場景都優於其餘框架，但Ray集合訓練、調參以及部署爲一體，仍能保持不錯的性能，於是值得學習和使用。

得益於Ray框架良好的性能，Ray普遍用於工業界（如螞蟻金服），要先學會使用Ray必先了解Ray的構成，下一小節就Ray的構成進行介紹。

2. Ray的使用

2.1 Ray的構成

Ray大體由四部分組成：

• Tune: 超參數調整模塊
• RLlib: 強化學習模塊
• RaySGD: 分佈式訓練模塊
• Ray Serve: 應用服務部署模塊

Ray涉及了AI應用的整個生命週期：訓練、調參、部署，並對強化學習場景進行了專門的優化。因爲我的使用經驗有限，這裏只介紹Ray的Serve模塊。

2.2 Ray的啓動

如圖4，Ray由一個頭節點（Head node）和一組工做節點（Worker node）組成。啓動Ray須要首先啓動頭節點，併爲工做節點提供頭節點的地址以造成集羣。頭節點負責管理和分配工做節點的任務，工做節點負責執行任務並返回結果。通過測試，頭節點和工做節點能夠爲同一臺計算機。

Ray的啓動由兩個步驟組成：啓動頭節點、註冊工做節點到頭節點。

圖4. Ray節點示意圖

如下是頭節點的啓動代碼和關閉代碼。

import ray
ray.init()  # 啓動
assert ray.is_initialized() == True

ray.shutdown()  # 關閉
assert ray.is_initialized() == False

注：啓動腳本應當加入關閉代碼，若是沒有，ray程序可能一直在進程中運行。

Ray框架採用Actor模型，相較於傳統的共享內存模型，Ray不存在狀態競爭、能夠方便的組建集羣、能更好的控制狀態。每一個Actor即每一個工做節點的註冊方式以下。

import ray
ray.init(address=頭節點地址)  # 啓動
assert ray.is_initialized() == True

ray.shutdown()  # 關閉
assert ray.is_initialized() == False

2.3 Ray Serve

Ray Serve能夠類比clipper，主要用於模型的部署服務，並支持多種深度學習框架，官方給出的示例有：

• Keras and Tensorflow Tutorial
• PyTorch Tutorial
• Scikit-Learn Tutorial

這裏以tensorflow2爲例，來講一下如何用ray來部署模型服務。

步驟一：定義一個模型服務類

以下是模型服務類的簡易代碼，和Flask等框架部署AI服務相似。因爲Ray使用gRPC做爲通訊協議，速度更快，Ray還在gRPC基礎上進行了優化，有些場景快於原生的gRPC通訊。

class TFMnistModel:
    def __init__(self, model_path):
        import tensorflow as tf
        self.model_path = model_path
        # 加載模型
        self.model = tf.keras.models.load_model(model_path)

    async def __call__(self, starlette_request):  # 異步調用
        # transform HTTP request -> tensorflow input
        input_array = np.array((await starlette_request.json())["array"])
        reshaped_array = input_array.reshape((1, 28, 28))

        #  tensorflow input -> tensorflow output
        prediction = self.model(reshaped_array)

        # 返回結果
        #  tensorflow output -> web output  
        return {
            "prediction": prediction.numpy().tolist(),
            "file": self.model_path
        }

步驟二：模型部署到Ray Serve

以下代碼中，start函數用於啓動服務，create_backend函數用於啓動模型，create_endpoint函數啓動服務。在Ray中，模型和服務是分離的，能夠多個服務調用同一個模型，以支持複雜的調用邏輯。

"tf:v1"爲模型的名稱，"tf_classifier"爲服務的名稱，route參數爲路由，這些參數均可自由定義。

client = serve.start()
client.create_backend("tf:v1", TFMnistModel, TRAINED_MODEL_PATH)
client.create_endpoint("tf_classifier", backend="tf:v1", route="/mnist")

步驟三：請求測試

resp = requests.get(
    "http://localhost:8000/mnist",
    json={"array": np.random.randn(28 * 28).tolist()})
print(resp.json())

3. 結語

一個優秀的框架每每包含了衆多先進的設計理念。Ray框架在構建時，參考了許多先進的設計理念，如混合調度策略、GCS 管理等等，這些設計理念使得框架自己完善而又先進。Ray普遍用於AI企業的分佈式計算場景，從衆多框架中脫穎而出，值得學習。