KubeCon North America 2024からAIワークロードのスケジューリングに関するセッションを紹介

KubeCon＋CloudNativeCon North America 2024から、前日のCloud Native AI Dayで行われたセッション「Multitenancy and Fairness at Scale with Kueue: Case Study」を紹介する。これはGoogleとAppleのエンジニアが一緒になってプレゼンテーションを行ったもので、クラスター内で実行されるバッチジョブ、主に機械学習のデータセットを学習する際にバッチで投入されるジョブをどのようにノードに配分するのか？ について、KueueというKubernetesに実装したオープンソースのスケジューラーを中心に解説する内容となっている。

クラスター内のバッチジョブをスケジューリングするKueueの解説セッション

動画は以下のリンクから参照して欲しい。

●動画：Multitenancy and Fairness at Scale with Kueue: A Case Study

セッションを行うGoogleとAppleのエンジニア

プレゼンテーションを行ったのはGoogleのAldo Culquicondor氏とAppleのRajat Phull氏だ。GoogleとAppleのエンジニアが一緒になってプレゼンテーションを行うというのも珍しい。さらに、KueueがKubernetesのワーキンググループ（バッチジョブのSIG）の配下で開発されていることからGoogleとAppleだけではなく、Red Hatなど多数の企業のエンジニアが結集して開発を行っている点も興味深い。

Kueueとは何かを解説

ここではKueueについて概要を説明し、Kubernetesのスケジューラー（kube-scheduler）とオートスケーラー（cluster-autoscaler）と連携することで、Kubernetes上でバッチジョブを効率的に配備することを可能にするソフトウェアであると解説した。マルチテナントで運用されるという前提の元、テナントごとに割り当てられたクォータ、テナント間でのリソースの貸し借りの設定、階層的な組織設定によるスケジューリングなどを可能にする点がスライドで解説されている。またあくまでもKubernetesが前提となるとしてPodやJob、JobSetなどのKubernetesのオブジェクトを利用し、KubeFlowやKubeRayなどとも連携が可能であるという。

クラスターの有限なリソースをどうやって割り当てるのか？ 問題の解決方法

ここからは複数のデベロッパーがバッチジョブを投入するという例を挙げて、リソースの割り当ての方法を解説する流れとなった。

A、B、Cのそれぞれのデベロッパーがジョブを投入

クラスター、特にGPUを実装したサーバーにどうやってジョブを配備するのが最適なのか？ という問題は、多分にそのジョブが構成されている文脈によって異なると思われる。Gang-Schedulingは複数のジョブを複数のノードで一斉に配備する方法で「All or Nothing」という発想に基づいてジョブが要求するリソースがなければスケジューリングを行わず、リソースが空いた状態になって初めて複数のジョブを実行するという方法になる。

その際に複数のデベロッパーが要求するリソースの総量が現実のクラスターのリソースを超える場合、いかにリソースを活用するのか？ という問題が発生する。それを解決するための発想がFair Sharingだ。

複数のデベロッパー間でリソースを配分するFair Sharingの解説

この例では緑色のデベロッパーがジョブを投入していない状態で、そのデベロッパーが持つクォータから黄色と青色のデベロッパーがリソースを借りてジョブを実行することを説明。それを決定するためにShareValueと呼ばれる数値が使われる。これはデベロッパーが持つクォータに重みを掛けたものを分母に、現状で借りているクォータを分子にした分数値となり、設定されたリソースに対して現在利用しているクォータが少ないジョブにリソースを多く割り当てるという方法だ。

クラスターの利用者を階層的な組織として定義し、スケジューリングに利用

ここではクラスターの利用について組織を階層化することでその優先順位を設定しようという機能を説明した。

階層的な組織を設定してどのジョブを優先的に実行するのかを決定

ここではジョブに関して「組織」という属性を定義し、その組織に属する下部組織は上位の組織の優先順位やクォータを継承することで、ジョブ全体が破綻なくスケジューリングされることを目指す機能となる。この優先順位算出もShareValueを使ったFair Sharingとなる。この例では中央のBack Propagationというグループが最も割り当てクォータが少ないため、優先的にリソースを獲得できることになる。

Kueueを使わないスケジューリングとKueueを使ったスケジューリングを比較

ここではKueueを使わない場合に比べて、クォータの管理とより詳細な第2レベルのスケジューリングが可能になると説明。ここで「第2レベル」と表現しているのはKubernetesにおいてはkube-schedulerとcluster-autoscalerが上位のスケジュールを受け持ち、その下でより粒度の細かなスケジューリングを担当するという意味だろう。実際にはKarpenterというKubernetesに特化したスケジューラーもAWSによって開発されており、スケジュールされていないPodに対するノードの割り当てを動的に行うことも可能となっている。Kueueはジョブをグループ化するCohortという属性や階層的組織という新たな属性をKueueに追加することで、機械学習などのジョブに合ったスケジューリングを行うという部分が差別化のポイントだろう。

なおKarpenterについては以下の公式サイトを参照されたい。

●参考：Karpenter

機械学習ジョブの割り当ての例

このスライドでは4つのテナント（ここでは一般的なマルチテナンシーにおける別組織、別会計を目指したものではなく、企業内の組織をひとつのテナントとして想定したものと思われる）に対してジョブの実行と停止（Preemption）がどのような戦略で行われるかを解説。最初に出てくるのはGang-Schedulingであるという部分にこの例が機械学習、つまり人工知能のための割り当て方法であり、KubeCon共催のCloud Native AI Dayで行われている理由でもあると言えるだろう。

3番目には階層的組織に基づくスケジューリングが記載されているが、実際のところこれはKueueのv0.9で実装予定の新機能であり、ここでの解説は多分に未来を先取りしているとも言えるし、実際にいつ頃実装されるのかは未定であることは注意しておくべきだろう。

特定のテナントがリソースを占有することを防ぐためのスケジューリング

ここでは使われていないリソースに対して、特定のテナントのジョブが占有してしまうことを防ぐための仕組みを解説。重みづけを行ったFair Sharingを行うことで公平なリソース管理が可能になることを説明した。

階層的組織を使ったリソース管理の例

組織を設定することでその組織の特性に応じたスケジューリングが可能になることも説明。ここも新機能として実装される予定の内容となる。

最後のスライドではバージョン番号で言えば、0.7、0.8、0.9という3つのリリースにそれぞれの新機能が実装される予定であることを説明した。またクラスター管理者向けのコマンドラインツール、kueuectlも予定されているという。1枚のスライドでは説明できないほどの多数の新機能が予定されていることが見てとれる。クラスターを複数使ってマルチクラスターで分散ジョブを管理するSkyRayや今回紹介したKueueなど、機械学習から生成型AIに特化したミドルウェア、管理ツールがこれからは多数発表されていくだろう。先端的な生成型AIのデベロッパーを抱えるユーザー企業においてはどの時点でどのツールに賭けるのかはギャンブルかもしれない。生成型AIのモデル以上にこれからは企業による「ツールの評価力」が試されていくことは間違いない。