9月1日、Collabnixが「Kueue on Kubernetes: GPU Job Queueing and Fair-Share Quotas for AI Workloads」と題した記事を公開した。この記事では、KubernetesのGPUジョブ管理における「デッドロック」問題をKueueで回避し、チーム間のフェアシェアクォータを実現するハンズオン手順について詳しく紹介されている。
GPUが8枚あるのに誰も使えないという問題
AIワークロードを動かしているKubernetesクラスターで起きがちな状況がある。GPUが8枚あり、3チームで共有している。あるチームが6枚を要求するジョブを、別のチームが4枚を要求するジョブを同時に投入すると、KubernetesのデフォルトスケジューラーはそれぞれのPodを順番にスケジュールしてしまう。結果、両ジョブが中途半端にGPUを確保したまま、どちらも動き出せずにハードウェアが遊ぶ。
デフォルトスケジューラーに欠陥があるわけではない。スケジューラーはPodを一つずつスケジュールするもので、「このジョブは全Podが揃わないと意味がない」という概念を持たない。また「researchチームにGPU 4枚、inferenceチームに4枚」という約束も知らない。
**Kueue**(Kubernetes SIGプロジェクト)はこの問題を解決するために設計されている。スケジューラーの上位に位置し、定義したクォータに基づいてどのジョブを起動するかを決定する。必要リソースが確保できないジョブはPodを一切作らずキューで待機する。これがデッドロックを防ぐ核心だ。
Kueueの4つの概念
Kueueを使いこなすには4つのオブジェクトを理解する必要がある。
- ResourceFlavor:ハードウェアの種類を表す。A100ノード用、L4ノード用、CPU用といったフレーバーを定義し、ノードラベルやtaintに対応づける
- ClusterQueue:クォータを保持するクラスタースコープのオブジェクト。プラットフォームチームが管理し、各フレーバーに対して利用可能なリソース量を定義する
- LocalQueue:Namespace内に存在し、ClusterQueueを参照する。アプリチームはLocalQueueにジョブを投入するだけで、ClusterQueueには触れない
- Workload:KueueがJobごとに自動生成するオブジェクト。ジョブが詰まっているときの診断に使う
この設計により、プラットフォームチームが一箇所でキャパシティを管理し、各チームはキュー名をマニフェストに書くだけでよい。クラスター全体の権限を持たなくてもトレーニングジョブを投入できる。
ハンズオン:2チームのGPUクォータを設定する
記事ではGPUなしのローカルクラスター(Kind)で動作確認できるハンズオンが紹介されている。CPUクォータで挙動を確認し、後半でGPU向けの変更点を解説する構成だ。
インストール
kind create cluster --name kueue-lab
kubectl apply --server-side -f \
https://github.com/kubernetes-sigs/kueue/releases/download/v0.18.2/manifests.yaml
kubectl -n kueue-system wait --for=condition=Available \
deployment/kueue-controller-manager --timeout=300s
※ 上記コマンド中のv0.18.2は元記事執筆時点のバージョンである。最新バージョンはKueueのリリースページで確認していただきたい。
APIバージョンの確認も忘れずに。Kueueはv1beta1からv1beta2に移行済みだ。
kubectl api-resources | grep kueue
ClusterQueueの設計:cohortとborrowingLimitが重要
2チーム分のClusterQueueを作成し、同じcohort(ai-platform)に所属させる。これにより、一方のチームが使っていないクォータを他方が借用できるようになる。
apiVersion: kueue.x-k8s.io/v1beta2
kind: ClusterQueue
metadata:
name: research-cq
spec:
namespaceSelector: {}
cohort: ai-platform
queueingStrategy: BestEffortFIFO
resourceGroups:
- coveredResources: ["cpu", "memory"]
flavors:
- name: default-flavor
resources:
- name: cpu
nominalQuota: 4
borrowingLimit: 4
- name: memory
nominalQuota: 4Gi
borrowingLimit: 4Gi
3つのフィールドが肝だ。
- **
nominalQuota**:そのチームに保証されたリソース量 - **
cohort**:同じcohort内で未使用クォータを共有できるようにする - **
borrowingLimit**:借用量の上限。これがないと、一方のチームがアイドル中に他方がリソースを全部使い切り、本来保証されているはずのチームが長時間待たされる
queueingStrategyのBestEffortFIFO(デフォルト)は、大きなジョブが詰まっていても小さなジョブが先にスケジュールされるため、クラスターの稼働率を高く保てる。StrictFIFOは厳密な順序を守るが、キューの先頭が通過するまで後続が全て待機するため、キャパシティが遊びやすい。
ジョブの投入方法
Kueueへの対応に必要なのはラベル1行だけだ。
labels:
kueue.x-k8s.io/queue-name: team-queue
そしてsuspend: trueを明示する。Kueueがクォータを確認し、空きがあればsuspend: falseに切り替えてジョブを起動する。クォータが足りなければPodは一切作られず、Workloadオブジェクトにその理由が記録される。
クォータが埋まった状態で追加ジョブを投入すると、実行中のジョブが完了した数秒後に自動的に次のジョブが起動する。手動介入もcronジョブもリトライループも不要だ。
借用とプリエンプション:安全に使うための設定
inferenceチームのクォータが空いていればresearchチームが借用できるが、後からinferenceチームがジョブを投入した場合のデフォルト動作は「researchのジョブが終わるまで待つ」だ。これを変えるにはプリエンプションを設定する。
kubectl patch clusterqueue inference-cq --type=merge -p '
spec:
preemption:
reclaimWithinCohort: Any
withinClusterQueue: LowerPriority'
reclaimWithinCohort: Anyを設定すると、inferenceチームは借用中のジョブを強制排除して自分のnominalQuotaを取り戻せる。ただしプリエンプションはジョブの強制終了を意味するため、チェックポイント機能なしのトレーニングジョブには設定しないこと。8時間分の計算が無駄になる。
実GPUへの対応
変更点は小さい。ResourceFlavorにGPUノードのラベルとtolerationを追加し、ClusterQueueのcoveredResourcesにnvidia.com/gpuを加えるだけだ。
apiVersion: kueue.x-k8s.io/v1beta2
kind: ResourceFlavor
metadata:
name: a100-flavor
spec:
nodeLabels:
nvidia.com/gpu.product: NVIDIA-A100-SXM4-80GB
tolerations:
- key: nvidia.com/gpu
operator: Exists
effect: NoSchedule
GPUモデルごとに別フレーバーを作成することが推奨されている。80GBのGPUメモリを必要とするジョブが、L4のクォータで受理されてしまうのを防ぐためだ。元記事ではKubernetes 1.34以降でDynamic Resource Allocation(DRA)を使用している場合、最新のKueueはDRAクレームもクォータに計上できると紹介されている(バージョン情報は元記事に基づく)。
バッチJobだけではない
KueueはKubernetesのbatch/v1 Jobだけでなく、JobSet、Kubeflowのトレーニングオペレーター(PyTorchJob、TFJob)、RayJob、RayCluster、MPIJob、さらにはPodやDeploymentにも対応している。いずれもkueue.x-k8s.io/queue-nameラベルを追加するだけで統合できる。つまり、既存のワークロード定義をほぼ変えることなく、クラスター全体のリソース配分をKueueの管理下に置けるということだ。GPUクォータの整備に取り組むチームにとって、導入コストの低さは大きな利点といえる。
詳細はKueue on Kubernetes: GPU Job Queueing and Fair-Share Quotas for AI Workloadsを参照していただきたい。