KAI Schedulerは、AIワークロード（AI関連のタスク）の効率的な管理を目的としたツール「NVIDIA Run:AI」の中核機能として開発された。GPU不足が深刻化する中で、限られたGPUリソースを複数のチームやプロジェクト間で公平かつ効率的に分配するための仕組みを提供する。NVIDIAによれば、KAI SchedulerはGPUリソースの需要変動を管理し、割り当てまでの待機時間を短縮する。GPUリソースの保証や適切な割り当ても実現している。

　GitHubの公式リポジトリの情報によると、KAI Schedulerは、小規模な対話型タスクから大規模なトレーニングや推論まで、Kubernetesクラスタ内におけるさまざまなAIワークロードを対象にする。NVIDIAは、異なるアプリケーション間での公平性を維持しつつ、最適なGPUリソース割り当てを保証すると説明している。

　このツールにより、Kubernetesクラスタの管理者は、GPUリソースを動的にワークロードに割り当てられるようになる。KAI Schedulerは他のスケジューラーと共存して動作することも可能だ。

　NVIDIAのソフトウェアシステム担当バイスプレジデントのロネン・ダール氏と、NVIDIA Run:aiのデータサイエンティストであるエキン・カラブルト氏は、ブログ記事で次のように述べている。「例えば、データ探索のような対話型作業にはGPUが1つあれば足りる一方で、分散トレーニングや複数の実験を実施する場合には、複数のGPUが必要なことがある」。従来のスケジューラーは、こうしたニーズの変動に対応できないという。

　KAI Schedulerは、公平なリソースシェア値を継続的に再計算し、クォータ（割り当て上限）や制限値をリアルタイムで調整する。こうした動的アプローチにより、管理者が手動で介入する必要なく、効率的なGPUリソース配分が可能になるという。

　機械学習エンジニアは、KAI Schedulerを以下のような手法と組み合わせることで、ジョブの待機時間を短縮する。

• ギャングスケジューリング
  • 複数のプロセッサを集団として扱い、同期させてスケジュールする手法。
• GPU共有
  • 1つのGPUを複数のジョブで時間的・論理的に分割して使えるようにする技術。
• 階層型キューイング
  • 複数のリクエストを階層構造の待ち行列に入れて、優先度を管理する方式。

　これにより、ユーザーはジョブをバッチとしてまとめて投入でき、GPUリソースが利用可能になり次第、優先度や公平性に基づいて順次実行される。

　GPUリソースの需要変動に対応するために、KAI Schedulerは「ビンパッキング」を使用している。ビンパッキングとは、限られた容量の容器であるビンに、さまざまなサイズの物体を効率よく詰め込むアルゴリズムだ。未使用のGPUリソースに小規模タスクを優先的に割り当てることで、リソースの断片化を防ぎ、計算資源の利用率を最大化する。

　ノード（コンテナを実行するための物理サーバや仮想マシン）の断片化にも対処するため、KAI Schedulerはタスクをノード間で再割り当てする。ワークロードを分散させることで、各ノードの負荷を最小化し、リソースの可用性を高めている。

　共有クラスタの運用においては、一部ユーザーがGPUリソースを過剰に確保し、GPUの未使用クォータが発生するケースがある。「一部の研究者はGPUの利用を確保するために、必要以上のGPUを早めに確保することがある。これにより、他のチームに未使用クォータがあるにもかかわらず、リソースの未活用を引き起こすことがある」とダール氏とカラブルト氏は説明する。KAI Schedulerはリソース保証機能を通じて、GPUリソースの独占を防ぎ、Kubernetesクラスタ全体の効率性を高めるという。

　KAI Schedulerは、Pod（Kubernetesでの基本的な実行単位）のグループを自動的に検出および管理する機能を備えている。これにより、以下のような機械学習フレームワークやツールを自動検出して接続する。

• Kubernetes向け機械学習ツールキット「Kubeflow」
• プログラミング言語「Python」での並列分散処理をシンプルかつ高速にするためのフレームワーク「Ray」
• Kubernetes向けの継続的デリバリーツール「Argo CD」

　これにより、Kubernetesクラスタの構成の複雑さを軽減し、開発スピードを向上させるとNVIDIAは主張している。

（翻訳・編集協力：編集プロダクション雨輝）