📝

Kubernetes でAIハードウェアの性能を最大限に引き出すためのスケジューラとデバイスプラグイン (1) 理解編

2024/05/17に公開

Kubernetes

OpenShift

tech

はじめに

大規模言語モデル (LLM) や基盤モデル (FM) などの学習、それらのモデルを処理対象に適用する推論、といった AI 関連の計算は、計算量の大きさ・メモリ使用量の多さ・分岐が少なくひたすらに数値計算を行う、といった特徴から、CPU 以外の専用デバイス（GPUなど）を使うことが事実上標準となっています。また、大規模な計算を複数計算機でスケールアウトさせるためには、計算機の効率を最大限に引き出せて、かつユーザが容易にジョブを定義できる環境が不可欠であり、そのための基盤として欠かせない存在になっているのが Kubernetes です。

この記事では、Kubernetes でAIハードウェアの性能を最大限に引き出すために必要な知識とソフトウェアについて解説します。想定する読者は以下のような皆さんです。

Kubernetes クラスタ管理者: Kubernetes に関する知識は持っているが、AIワークロードの実行には携わったことがないので、AI向けハードウェアを搭載したノードを提供するにあたってどのような準備が必要で、どのような点に注意すればよいのか知りたい
AI エンジニア: これまでAI向けハードウェアを利用したプログラムをスタンドアロンで実行した経験はあるが、それを Kubernetes で実行するにはどうすればよいのか知りたい
AI ハードウェア用ソフトウェア技術者: 自組織が開発したハードウェアを Kubernetes で使ってもらえるようにソフトウェア群を実装したい、もしくは自組織で使っているAIハードウェアを極限まで効率よく使いたいからソフトウェア群も自前で実装したい

全ての内容をひとつの記事にするには長くなり過ぎるので、以下の2編に分割して記述していきます。

理解編（この記事）: GPUを始めとするAI向けハードウェアを Kubernetes で動く AI ワークロードから使うためにはどのようなソフトウェア群が必要で、それはどのような点に気をつけて使うべきか、現状の課題と将来の見込みを説明します。
設計・実装編（次回の記事）: 理解編で解説した課題を克服する方法について検討し、AI アクセラレータ向けの具体的な実装例を紹介します。こちらの記事はAI ハードウェア用ソフトウェア技術者の方向けに特化しているので、Kubernetes管理者の方やAIエンジニアの方は読み飛ばしても大丈夫です。

それではまず「理解編」です。

Kubernetes で CPU 以外のリソースをリクエストしたときの動き

Kubernetes クラスタで何らかの計算を行いたいとき、ユーザは計算処理をコンテナの中に記述して、Pod という形でクラスタにデプロイします。具体的には、下記のような YAML ファイルを記述した上で、これを Kubernetes API サーバに POST します。

apiVersion: v1
kind: Pod
metadata:
  name: my-ai-app-1
spec:
  containers:
    - name: my-ai-app-1
      image: "registry.example.com/my-ai-app:1.0.1"

ここで、このコンテナがCPUやメモリ以外のリソース、例えば NVIDIA (R) CUDA (R) を経由して GPU を利用するプログラムを含んでいるとすると、YAML には GPU というリソースをリクエストする記述を追加する必要があります。

apiVersion: v1
kind: Pod
metadata:
  name: my-ai-app-1
spec:
  containers:
    - name: my-ai-app-1
      image: "registry.example.com/my-ai-app:1.0.1"
      resources:
        limits:
          nvidia.com/gpu: 1 # requesting 1 GPU

nvidia.com/gpu: 1 というリクエストを Kubernetes に解釈してもらうには、GPU を搭載したノードを用意した上で、NVIDIA が提供する Device Plugin を Kubernetes に導入する必要があります。

この manifest が Kubernetes API Server に到着したあと、デバイスはどのように割り当てられるのでしょうか？　まず最初に Kubernetes が考えなければならないのは、この Pod をどのノードに割り当てようか、ということです。CPUやメモリに余裕があるノードを選択する必要があるのはもちろんですが、manifest を見るとGPU をリクエストしているのですから、 GPU が搭載されていて、かつ未使用のGPUがリクエストされた量以上に残っているノードを選択的に割り当てなければいけません。

ノードの割り当ては Kubernetes 内部のコンポーネントのうち、Scheduler と呼ばれるコンポーネントの仕事です。Scheduler は Node リソースを参照して、要求されたリソースが要求された量だけ割り当て可能なノードを探して、割り当てるノードを決定します（ドキュメント）。

続いて manifest は割り当てられたノードの Kubelet に引き渡されます。Kubelet はリソース名に対応する Device Plugin の助けを借りつつ、指定された Pod 内のコンテナにCPU・メモリ・デバイスといったリソースを割り当てた上でコンテナを起動します。このとき、 Kubelet は Node リソースを更新し、割り当て可能なデバイスが少なくなったことを登録しておきます。同様に、Pod 停止時は解放されたデバイスを割り当て可能なデバイスとして Node リソースに登録し直します。これにより、スケジューラは Node リソースを参照することで常に最新のデバイス割り当て状況を把握できます。

ここでもう一度、ユーザ視点でさきほどの YAML ファイルを見返してみます。ユーザが指定しているのは、リソース名とその量だけです。どのノードのどのデバイスを使うかは指定していません。

...
  resources:
    limits:
      nvidia.com/gpu: 1 # requesting 1 GPU

つまり、

どのノードにどれくらいの量のデバイスが割り当て可能状態で残っているのか
割り当て可能なデバイスのアドレスは何なのか

といった、インフラの状況に依存する事項を気にする必要はありません。言い換えると、Kubernetes はアプリケーションに対してインフラ部分を隠蔽してくれている、ということになります。クラスタ上でAIワークロードを実行する利用者にとっては、「（リソース名という形で表現された）デバイスの種類」及び「必要な量」だけを指定すればよいので、 manifest の記述が容易になり、使い回しもしやすいというメリットがあります。

では、このしくみはAIハードウェアを使う上で完全無欠なのかというと、そうとも言えないのです。なぜ「そうとも言えない」のかを、次節で説明します。

AIワークロードと Kubernetes の摩擦点

さきほどの動作例では、AIワークロードはひとつのコンテナ・ひとつの Pod で動くことを想定していました。しかしAIワークロードはそんな量の計算リソースで済むような計算量ではないのがふつうです。複数のAIハードウェア搭載した複数の計算機で動かし、それぞれのデバイスで更新したパラメタを相互に転送して次の iteration を実行します。となると、デバイスやノードのトポロジーが学習・推論の性能に影響を与える、ということになります。

例えば、以下のようなトポロジーを持つ2つの計算機ノードを仮定しましょう。計算機内では PCI Express バス上に GPU や NIC といったデバイスが接続され、NIC を介して相手の計算機と接続されています。

ここでもし、AIワークロードがこの2ノード両方で実行され、それぞれのプロセスが赤色で示されたGPUの割り当てを受けたとするとどうなるでしょうか？　パラメタ更新やモデル転送といったやりとりが発生した際、ノード内・ノード相互でのデータ転送時に PCI Bridge を跨いだ通信となり、RDMA (remote direct memory access) などの高速な転送手段が利用できなくなる可能性があります。となると、リソースをリクエストする際に「ネットワーク視点で近傍にあるデバイス（上図で言うと青色のGPU）を割り当てて欲しい」といった細かな条件をつけたくなります。

ところが前節で触れたように、Kubernetes はインフラ部分を隠蔽する形でリソースを提供します。そして、その思想は内部設計にも反映されています。

デバイス割り当てを担う Device Plugin が Kubelet からデバイス割り当てをリクエストされる際に呼ばれる gRPC サービス Allocate() を見てみます。ここで Kubelet から Device Plugin に渡されるのは AllocateRequest というメッセージです（ドキュメント）。

service DevicePlugin {
...
      rpc Allocate(AllocateRequest) returns (AllocateResponse) {}
}

AllocateRequest は Go の実装としては ContainerAllocateRequest のスライス（へのポインタ）からなり、その中身はデバイスIDのスライスです。

type AllocateRequest struct {
	ContainerRequests    []*ContainerAllocateRequest `protobuf:"bytes,1,rep,name=container_requests,json=containerRequests,proto3" json:"container_requests,omitempty"`
	XXX_NoUnkeyedLiteral struct{}                    `json:"-"`
	XXX_sizecache        int32                       `json:"-"`
}
...
type ContainerAllocateRequest struct {
	DevicesIDs           []string `protobuf:"bytes,1,rep,name=devices_ids,json=devicesIds,proto3" json:"devices_ids,omitempty"`
	XXX_NoUnkeyedLiteral struct{} `json:"-"`
	XXX_sizecache        int32    `json:"-"`
}

ここで、渡されるデバイスID（これから割り当てようと思っているデバイスのリスト）を決めているのは Kubelet です。実は、デバイスの割り当て状況を管理しているのは Kubelet だけで、Device Plugin は割り当て状況を管理していません。 基本的に、「どのデバイスが未割り当て／割り当て済みで、次に来たリクエストに対してどのデバイスを割り当てるか」を決めるのは Kubelet の仕事です。Device Plugin にリクエストが来た時点で、割り当てるデバイスは Kubelet によって決められたあとです^[1]。

また、Device Plugin からは Pod manifest も見えません。 ですので、例えば「manifest に annotation を付けてデバイス割り当ての細かいリクエストを Device Plugin に伝えたい」と思っても、Device Plugin のしくみでは、それを実現する手だてはないのです。

DRA (Dynamic Resource Allocation)

Device Plugin のしくみでできないのはわかったけど、でもAIハードウェアを使うとしたらやはり細かいデバイス指定は必須だろう…ということで、Kubernetes SIG Node で策定作業が進められているのが DRA (Dynamic Resource Allocation) と呼ばれる新しい feature です。

先に書いておくと、2024年5月現在、この feature の成熟度は alpha であり、production 環境では使えません（例えば商用 Kubernetes である OpenShift でこの機能を使うとサポート対象外となります)。また、DRA は更なるアーキテクチャの変更が予定されており、こちらの記事によると Kubernetes 1.32 での beta 化を目指しているとのことです。Kubernetes 1.32 はまだリリース予定が決まっておらず、Kubernetes 1.31 のリリースは2024年8月に予定されていますから、ここ数ヶ月は DRA を実用途で使える目処はない、と思った方がよさそうです。

DRA についてはこちらの日本語記事がとても参考になります。端的に言うと DRA とは、Kubernetes に既に存在している Storage Class / Persistent Volume / Persistent Volume Claim のようなしくみをリソースでも実現するためのフレームワークだと言えます。Resource Claim という形で必要なリソースを宣言しておいて、Pod manifest にはどの Resource Claim を使いたいのかを記述します。DRA ドライバは PodSchedulingContext というカスタムリソースを使ってスケジューラと協調しながら割り当て作業を行います。

DRA と Device Plugin の主な相違点は以下の通りです。

Pod manifest から渡せるパラメタ

前述の通り、Pod manifest から Device Plugin に伝達できるのは必要なデバイスの個数だけです。DRA のしくみの下では、Pod manifest から Resource Claim もしくは ConfigMap を経由して DRA ドライバにパラメタを渡すことができます。渡せるパラメタは、DRA ドライバが予め定義しているパラメタです (NVIDIA k8s-dra-driver の例)。

Device Plugin / DRA ドライバが呼び出されるタイミング

Device Plugin が Kubelet から呼び出されるのはデバイス割り当て時だけで、デバイスが解放されても Device Plugin には何らの呼び出しも行われません。このため、デバイス解放時の後始末を実装することが困難です。一方、DRA Driver は解放時にも呼び出しが行われるので　(NodeUnprepareResources) 、ここで解放処理を実装することができます。

デバイス割り当て状況の管理

Device Plugin はデバイス割り当て状況を管理しません。どのデバイスがどの Pod (コンテナ) に割り当てられているのかを管理するのは Kubelet の役割で、Device Plugin は PodResourcesLister　gRPC サービスを使って能動的に Kubelet に問い合わせない限り、どのデバイスがどの Pod (コンテナ) に割り当てられているのかを知る術がありません。DRA の場合、DRA ドライバは自身が割り当て・割り当て状況を管理します。

「理解編」まとめ

ここまでの記事をざっくり要約すると、以下のようになります。

AI ワークロード実行のために Kubernetes から AI ハードウェアを使いたいとき、現時点で商用利用可能なしくみは Device Plugin のみ（将来的には DRA が利用可能）
Device Plugin ではインフラの詳細部分が隠蔽されるため、AI ワークロード実行時に必要な細かなデバイスの指定を行うことが困難

このような状況で、「AI ワークロード実行時に必要な細かなデバイスの指定」を実現するにはどうしたらいいのか？　それを次の記事で解説したいと思います。

脚注

実は特別なやりかたとして、「Kubelet から Device Plugin に割り当て可能な全デバイスのリストを伝えて、Device Plugin がその中から好みのデバイスを選び出して Kubelet に返す」という割り当て方法もあります。実装編ではこの方法を使います。 ↩︎