Redisのキースペース通知とリースを使用したGPUワーカープールのロードバランシング

GPUのプールはステートレスなWebサーバーのプールとは異なり、それと同じようにロードバランシングを行うと、所有する最も高価なハードウェアを無駄にすることになります。重い推論ジョブを実行しているGPUは、完全に占有されます。2つ目のジョブを送信すると、それは共有されず、キューに入れられます。そして、別のGPUがアイドル状態である間に、リクエストは他の誰かの数分間の作業の後ろで待つことになります。Webロードバランシングから反射的に使われるラウンドロビンと最小コネクションは、GPUにとって「ビジー」が何を意味するのかを理解していないため、ここではどちらも失敗します。この投稿では、GPUを予約するリースと、GPUが解放された瞬間に通知するキースペース通知という2つのRedisプリミティブから、ビジーアウェアなスケジューラを構築し、その設計を余儀なくさせた障害モードについて説明します。

なぜ通常のロードバランシングはGPUで失敗するのか

通常のロードバランサーの前提は、GPUの処理には当てはまりません。Webサーバーは多数の同時リクエストを処理し、それぞれのリクエストは安価で短時間です。そのため、リクエストをサーバー間に均等に分散させることが最適に近い状態となり、古い負荷情報が問題になることはほとんどありません。GPUは一度に1つの重いジョブを処理し、それぞれのジョブは高コストで長時間です。そのため、唯一重要なのは特定のGPUが今すぐ空いているかどうかであり、その情報が古いと、即時開始できるか数分間待たされるかの違いになります。

ラウンドロビンは使用状況を完全に無視し、ジョブNをまだ処理中のGPUに平気でジョブN+1を割り当てます。最小コネクション方式は近いですが、それでも間違いです。なぜなら、GPUの占有率ではなくコネクション数を数えるためです。そして、推論中のGPUは、完全に利用不可能な場合でもコネクションを全く保持していないことがあります。必要なのは、「この特定のGPUは空いているか」という情報を信頼できる唯一の情報源とし、その状態が変化した瞬間に更新され、さらに2つのリクエストが同じ空きGPUを同時に要求できないように強力な予約機能を持つスケジューラです。Redisは、その両方を安価に提供します。

リースが予約となる

中核となる状態はGPUごとのリースです。これは、GPUがビジー状態のときに存在し、空き状態になると消えるRedisキーです。GPUを要求することは、そのリースを取得することであり、その要求はアトミックでなければなりません。そうすることで、同じ空きGPUをめぐって競合する2つのリクエストが両方とも勝つことはありません。

// 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は、1つのコマンドでcompare-and-swapを行うもので、キーがまだ存在しない場合にのみキーを設定します。競合に勝ったGPUはリースを取得し、ジョブを開始します。それ以外のGPUは、okがfalseで返されるのを見て、別のGPUを試すために次に進みます。ジョブが終了すると、ワーカーはリースを削除し、GPUは再び空き状態になります。もしこれだけだった場合、空きGPUを見つけるために依然としてポーリングを行う必要があります。そしてポーリングは、まさに私たちが避けたいレイテンシーと無駄そのものです。2つ目のプリミティブが、そのポーリングをなくします。

キースペース通知は「空き」をイベントに変える

Redisは、キーが変更されるたびにイベントを公開できます。これにはキーの有効期限切れや削除も含まれます。その機能がキースペース通知であり、待機中のリクエストがタイマーで問い合わせるのではなく、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メッセージが1つ移動する時間で、ビジー状態から再割り当て状態に移行し、定常状態ではポーリングは行われません。

TTLが存在する理由である障害モード

ハッピーパスには穴があり、それは最も高価なリソースを遊休させてしまうため、最も重要なものです。あるワーカーがGPUを要求し、ジョブを開始した後にクラッシュした、あるいはネットワークが分断された、あるいはそのプロセスが実行途中で強制終了されたとします。そのワーカーはリースを削除する行には決して到達しません。リースのキーはRedisに永久に残り、そのGPUは永遠に使用中としてマークされ、もはや存在しないプロセスからキーを削除するものは何もないため、キースペースイベントが発行されることは決してありません。一度のクラッシュで、GPUがプールから永久に削除されてしまいます。

これが、リースがTTL、つまり上記のSetNX呼び出しにおけるleaseTTLを持つ理由です。TTLは、ワーカーが完了時に有効期限が切れるずっと前にリースを削除する、通常ケースのためのものではありません。それは異常ケースのための回復メカニズムです。ホルダーが死んだ場合、リースは自動的に期限切れになり、Redisは期限切れのキースペースイベントを発行し、GPUは人手を介さずに自動的にプールに再参加します。このリースは、単一ランナージョブに関する前述の議論と同じ意味での真のリースです。つまり、生きているホルダーが更新し、死んだホルダーが失う予約です。

つまり、長いジョブは、TTLが経過する前に、まさしくハートビートのようにリースを更新しなければなりません。さもなければ、回復メカニズムによって実行途中でGPUを失うことになります。TTLは、予想される最長ジョブよりも十分に長く設定するか、長さを制限できないジョブについては定期的にリースをハートビートしてください。トレードオフはよく知られたものです。短いTTLはクラッシュしたGPUを迅速に回復させますが、遅いジョブを奪ってしまうリスクがあります。長いTTLは遅いジョブにとっては安全ですが、クラッシュしたGPUをより長く遊休状態にしてしまいます。デフォルト値からではなく、ジョブの長さの分布から選択してください。

確保と待機の間の競合

ビジーアウェアスケジューラには、書きやすく再現が困難であるため指摘する価値のある、微妙な順序のバグがあります。あるリクエストがすべてのGPUをチェックし、それらがビジーであることを見つけ、その後、空きイベントフィードをサブスクライブすると想像してください。もしチェックとサブスクライブの間のごくわずかな隙間でGPUが解放されると、サブスクリプションが存在する前にイベントが発行され、リクエストはそのイベントを検知できず、数分後になるかもしれない次の空きイベントを待つことになります。GPUは空いており、リクエストは待機していますが、どちらも他方の存在を知りません。

修正方法は、順序を逆にして、サブスクライブした後に再チェックすることです。まず空きフィードをサブスクライブし、その後GPUの確保を試みます。確保に成功した場合は、サブスクライブを解除して実行します。すべてがビジーであるために失敗した場合は、すでにサブスクライブしているフィードで待機します。そして重要なのは、イベントが到着したときに、解放されたGPUが自分のものだと仮定するのではなく、再度確保を試みることです。なぜなら、別の待機者が先にそれを確保している可能性があるからです。サブスクライブ、次にチェック、そして待機と再試行です。その順序によって隙間が埋められ、ウェイクアップ時の再試行は、解放されたGPUが自分のところに到達する前に他の誰かによって確保されてしまうという、複数の待機者がいるケースを処理します。

陳腐化した状態と調整パス

TTLがあっても、分散状態は乖離していきます。ワーカーが一時的に分断され、まだGPUを保持していると思い込んでいる間に、そのリースが削除されてしまう可能性があります。サブスクライバーがイベント発行の瞬間に再接続していた場合、期限切れイベントが見逃される可能性があります。イベントのみに基づいて構築されたイベント駆動システムは、時間とともにこれらの小さな乖離を蓄積していきます。その解決策は定期的な調整です。つまり、遅い間隔で、各GPUワーカーが実際に何を行っているかを独立して問い合わせ、現実に合わせてRedisを修正します。イベントは、一般的なケースにおいてシステムを高速かつ最新の状態に保ちます。調整パスは、イベントが見逃した乖離を捉えることで、長期的にはシステムを正しい状態に保ちます。イベントによる高速パス、そして調整による正確性のバックストップは、あらゆるイベント駆動スケジューラにとって望ましい基本姿勢です。

よりシンプルな設計で十分な場合

この仕組みは、GPUが希少で高価なリソースであり、ジョブが十分に長く、ビジー状態のGPUにジョブをルーティングすることがコストのかかる間違いとなる場合に、その価値を発揮します。ジョブが短く安価な場合は、アイドル状態のワーカーがプルする単純なキューの方がシンプルで、同様に優れています。各ジョブにかかる時間がわずかであれば、少し不完全な割り当てのコストは無視できるためです。配置を処理してくれるマネージド推論サービスを利用している場合は、それを使い、これらすべてをスキップしてください。まさにこの問題を解決するためにお金を払っているからです。

リースと通知のスケジューラーは、次の3つのことが同時に当てはまる場合に構築してください：GPUが十分に高価でアイドル時間が実質的な損失となること、ジョブが十分に長く誤った割り当てがリクエストを意味のある時間滞留させること、そしてマネージドエンドポイントをレンタルするのではなく自分でプールを運用していること。これらの条件下では、アトミックな予約のためのRedisリースと即時のハンドオフのためのkeyspace notificationが、ポーリングなしで高価なハードウェアをビジー状態に保つスケジューラーを提供し、TTLと調整パスが、ワーカーが停止したりイベントが失われたりした場合にその整合性を保ちます。