使用 Redis 鍵空間通知與租約對 GPU 工作者池進行負載平衡

GPU 資源池並不像一個無狀態（stateless）的網路伺服器池，若像對待後者那樣對其進行負載平衡，將會浪費您所擁有最昂貴的硬體。一個正在執行繁重推論作業的 GPU 會被完全佔用。若傳送第二個作業給它，它不會共享資源，而是會將作業排入佇列；於是，一個請求就得在別人長達數分鐘的工作後面等待，而此時另一個 GPU 卻閒置著。循環分配（Round-robin）和最少連線（least-connections）這些源自網路負載平衡的直覺反應方法，在這裡都會失效，因為它們不知道「忙碌」對 GPU 而言意味著什麼。本篇文章將基於 Redis 的兩種原語，建構一個能感知忙碌狀態的排程器：一種用以預留 GPU 的租約（lease），以及一種能在 GPU 一釋出便立即通知的鍵空間通知（keyspace notification）；本文也將探討促使此設計誕生的各種失敗模式。

為何一般負載平衡在 GPU 上會失效

一般負載平衡器所基於的假設，並不適用於 GPU 的工作。網頁伺服器處理許多並行請求，每個請求成本低廉且耗時短暫，因此將請求均勻分佈到各個伺服器上已接近最佳化，而過時的負載狀況幾乎無關緊要。GPU 一次只處理一個繁重的工作，每個工作都成本高昂且耗時長，因此唯一重要的是某個特定的 GPU 當下是否閒置，而關於此狀態的過時資訊，將會造成「立即開始」與「長達數分鐘停滯」之間的差別。

Round-robin 完全忽略了忙碌狀態，並且會毫無懸念地將第 N+1 個任務指派給仍在處理第 N 個任務的 GPU。Least-connections 演算法比較接近，但仍然是錯的，因為它計算的是連線數，而非 GPU 的佔用率，而一個正在進行推論的 GPU 可能在完全無法使用的情況下，卻沒有任何連線。你需要的是一個排程器，其判斷的依據是「這個特定的 GPU 是否閒置」，並在狀態改變的瞬間立即更新，且帶有硬性預留機制，如此一來兩個請求便無法同時佔用同一個閒置的 GPU。Redis 能以低廉的成本提供這兩方面的功能。

租約即是預留

核心狀態是每個 GPU 的租約：一個 Redis 金鑰，在 GPU 忙碌時存在，在閒置時消失。宣告一個 GPU 就是取得其租約，而此宣告必須是原子操作，如此一來，兩個競爭同一個閒置 GPU 的請求才不會都獲勝。

// Claim GPU n only if its lease is absent. SET NX is atomic, so exactly one
// caller wins the race. The TTL is a safety net, covered below.
ok, err := rdb.SetNX(ctx, "gpu:lease:"+gpuID, workerID, leaseTTL).Result()
if err != nil {
	return "", err
}
if !ok {
	return "", errBusy // someone else holds this GPU
}

SetNX 是一個指令中的比較並交換：僅在金鑰不存在時才設定金鑰。獲勝的 GPU 取得租約並開始其工作。其他所有人都會看到 ok 回傳 false，然後繼續嘗試另一個 GPU。當工作完成時，工作單元會刪除租約，GPU 就再次閒置。如果僅止於此，您仍然需要輪詢以尋找閒置的 GPU，而輪詢正是我們想要避免的延遲與浪費的權衡。第二個原語移除了輪詢。

Keyspace notifications 將「空閒」轉變為一個事件

每當一個鍵發生變化，包括鍵過期或被刪除時，Redis 都可以發布一個事件。這個功能就是鍵空間通知 (keyspace notifications)，它讓等待中的請求能夠在 GPU 空閒出來的瞬間得知，而不用透過計時器來詢問。您只需在伺服器上為您關心的事件類別（通用鍵事件和過期事件）啟用一次，然後訂閱租約鍵的模式。

// Subscribe to deletions and expirations of GPU lease keys. Each message
// means a specific GPU just became free.
pubsub := rdb.PSubscribe(ctx,
	"__keyevent@0__:del",
	"__keyevent@0__:expired",
)
for msg := range pubsub.Channel() {
	freedGPU := parseGPUID(msg.Payload) // the lease key that just vanished
	scheduler.notifyFree(freedGPU)
}

現在，整個流程是端到端的事件驅動。一個請求會嘗試用 SetNX 來宣告佔用任何空閒的 GPU。如果所有的租約都已被持有，它會等待，而不是空轉。任何地方的任何工作完成並刪除其租約，或租約過期的那一刻，Redis 就會發布事件，一個等待中的請求會被喚醒，並嘗試宣告佔用新空閒出來的 GPU。昂貴的硬體從忙碌到重新分配，所花費的時間僅為一條 Redis 訊息的傳輸時間，在穩定狀態下沒有輪詢。

TTL 所針對的故障模式

在理想路徑中存在一個漏洞，而這個漏洞至關重要，因為它會擱置您最昂貴的資源。假設一個工作單元取得了一張 GPU、開始了一項工作，然後崩潰了，或者其網路發生分區，或者其行程在執行中被終止。它永遠不會執行到刪除租約的那行程式碼。租約金鑰會永遠留在 Redis 中，那張 GPU 會被永久標記為忙碌，而且永遠不會為它觸發任何鍵空間事件，因為一個已不存在的行程不可能去刪除金鑰。一次崩潰就會將一張 GPU 從資源池中永久移除。

這就是為什麼租約帶有存留時間 (TTL)，也就是上述 SetNX 呼叫中的 leaseTTL。TTL 並非為正常情況設計，在正常情況下，工作單元會在完成工作時，遠在租約到期前就刪除它。它是針對異常情況的恢復機制。如果持有者死掉，租約會自行到期，Redis 會觸發到期鍵空間事件，而 GPU 會在沒有人為干預的情況下自動重新加入資源池。這個租約的真正意義，與先前討論單一執行者工作時的租約相同：一個由存活的持有者續約、由死掉的持有者放棄的預留。

這意味著一個長時間的工作必須在 TTL 到期前更新其租約，就像心跳一樣，否則它將在執行中途因恢復機制而失去其 GPU。對於預期最長的工作，請將 TTL 設定得比它長得多；對於無法界定長度的工作，則應定期對租約進行心跳續約。這是一個常見的權衡：短的 TTL 可以快速恢復崩潰的 GPU，但有撤銷慢速工作的風險；長的 TTL 對慢速工作是安全的，但會讓崩潰的 GPU 擱置更長時間。應根據您的工作時長分佈來選擇 TTL，而不是使用預設值。

聲請與等待之間的競態

一個忙碌感知排程器有個微妙的順序錯誤，這點值得提出來，因為它很容易寫出來，但很難重現。想像一個請求檢查了所有 GPU，發現它們都在忙碌中，然後訂閱了可用事件通知。如果在檢查和訂閱之間的微小間隙中，有一個 GPU 變為可用，事件會在訂閱存在之前觸發，該請求永遠不會收到通知，然後它會等待下一個可能在幾分鐘後才發生的可用事件。GPU 是可用的，請求正在等待，而彼此都不知道對方的存在。

解決方法是反轉順序，並在訂閱後重新檢查。首先訂閱可用通知，然後嘗試聲請一個 GPU。如果聲請成功，就取消訂閱並執行。如果因為所有資源都忙碌而失敗，就等待你已經訂閱的通知，而且至關重要的是，當事件到達時，要再次嘗試聲請，而不是假設被釋出的 GPU 就是你的，因為另一個等待者可能已經搶先一步拿走了它。訂閱，然後檢查，然後等待並重試。這個順序彌補了那個間隙，而喚醒時重試的機制處理了有多個等待者的情況，也就是在你來得及取得之前，被釋出的 GPU 就被其他人聲請了。

過時狀態與協調階段

即使有 TTL，分散式狀態也會發生偏移。當其工作單元短暫分區時，租約可能已被刪除，但該工作單元仍認為自己持有 GPU。如果一個訂閱者在事件觸發的當下正在重新連線，那麼該到期事件可能會被錯過。僅基於事件建構的事件驅動系統，會隨著時間的推移累積這些微小的分歧。解方是定期進行協調：以一個較慢的間隔，獨立地詢問每個 GPU 工作單元它實際上在做什麼，並修正 Redis 以符合實際情況。在一般情況下，事件能保持系統的快速與即時。從長遠來看，協調階段透過捕捉事件錯過的偏移來保持其正確性。由事件構成的快速路徑，由協調構成的正確性後盾，是任何事件驅動排程器的一個良好預設態勢。

什麼時候簡單的設計就足夠了

當 GPU 是你稀缺且昂貴的資源，且工作時間長到將其路由到一個忙碌的 GPU 會是個代價高昂的錯誤時，這套機制就有了用武之地。如果你的工作既短暫又便宜，一個任何閒置工作者都可以從中提取工作的普通佇列會更簡單且效果一樣好，因為當每個工作都只需片刻時，一個稍微不完美的分配所產生的成本可以忽略不計。如果你有一個代管的推論服務為你處理放置問題，那就使用它並跳過所有這些，因為你付費正是為了讓這個問題得到解決。

當三件事同時成立時，才需要建立租賃與通知排程器：GPU 昂貴到閒置時間是真正的損失、工作時間長到錯誤的分配會讓一個請求擱置一段有意義的時間，以及你是自己運行資源池而不是租用代管的端點。在這些條件下，用於原子性預留的 Redis 租賃和用於即時交接的鍵空間通知，為你提供了一個無需輪詢即可讓昂貴硬體保持忙碌的排程器，而 TTL 加上協調過程則在工作者死亡和事件遺失時保持其誠實運作。