Time-Slicing GPUs を Kubernetes で利用する

はじめに

Kubernetes にて、1つのGPUを複数コンテナ (※ Pod内の複数コンテナ、複数のPodを指す) で使い倒したい。そんな時はありますでしょうか。
本記事では、NVIDIA/k8s-device-plugin の Time-Slicing GPU を使うまでの準備 (GPUを利用するアプリ・Kubernetes の観点より) についてお話しします。

Kubernetes の用語やマニフェスト管理のツールについてある程度知見があるという前提で進めさせていただきます。
Kubernetes の文脈でのタイムスライシングは Time-Slicing GPUs、NVIDIA の文脈の場合は Timeslice と表記します。

NVIDIA/k8s-device-plugin の概要

DaemonSet で、以下の 3つの役割を持ちます。 (README.md より参照)

• 各ノードの GPU 数を公開
• GPUの状態を監視
• GPU対応コンテナの実行

Helm Chart を覗くと以下のコンポーネントが依存しています。

• node-feature-discovery
  • ノードで利用可能なハードウェア機能を公開 (CPU/Memory/Kernel/etc)
• gpu-feature-discovery
  • ノードで利用可能なGPU の情報を公開 GPU を公開

Time-Slicing GPUs について

Kubernetes 目線の話

Kubernetes で GPU のリソース割り当てを設定する場合、一般的には 1 つの GPU を 1 つのコンテナが占有する排他的アクセスとなります。

↓ コンテナにGPUを割り当てるサンプル YAML ↓

Kubernetes で GPU のリソース割り当てを設定する場合、一般的には 1 つの GPU を 1 つのコンテナが占有する排他的アクセスとなります。

apiVersion: v1
kind: Pod
metadata:
  name: example-vector-add
spec:
  containers:
    - name: example-vector-add
      image: "registry.example/example-vector-add:v42"
      resources:
        limits:
          gpu-vendor.example/example-gpu: 1

ここで、Time-Slicing GPUs を有効することにより、1つのGPUを複数のコンテナで共用利用できるようになります。
設定後は、既存のGPUリソース名の接尾辞に .shared がつき それを指定して共有GPUとして利用できるようになります (※ リソース名を変更するオプションを入れた場合) 。

↓ コンテナに共有GPUを割り当てるサンプル YAML ↓

apiVersion: v1
kind: Pod
metadata:
  name: example-vector-add
spec:
  containers:
    - name: example-vector-add
      image: "registry.example/example-vector-add:v42"
      resources:
        limits:
          gpu-vendor.example/example-gpu.shared: 1

GPU 目線の話

NVIDIA の Pascal アーキテクチャ以降では Compute Preemption がサポートされています。
Compute Preemption は、GPU上で実行中の計算タスクを命令のレベルで中断できます。この際、中断されたタスクの実行コンテキスト（レジスタや共有メモリなど）はGPU DRAMにスワップアウトされるため、他のアプリケーションがスワップインされて実行できるようになります。

Tuning CUDA Applications for Pascal, CUDA, リンク
Improving GPU Utilization in Kubernetes, NVIDIA, リンク

タイム スライシングは、複数のプロセスを同じ GPU 上でスケジュールできるようにするメカニズムです。スケジューラはすべての GPU プロセスに均等な時間を割り当て、ラウンドロビン方式で交互に実行します。

Efficient Access to Shared GPU Resources: Part 1, kubernetes@CERN, リンク

つまり、Time-Slicing GPUs (Timeslice)とは NVIDIA GPU アーキテクチャの機能であり、NVIDIA/k8s-device-plugin は、1つのGPU を複数のコンテナで利用できるように Kubernetes を調整する役割を持っています

利用する際の注意点としては

1. コンテナに 2 つ以上の Time-Slicing GPUs のリソースを割り当てたとしても、比例した量のGPU計算能力を利用できることは保証しません。
2. メモリ制限は Timeslice GPU には適用されません。1つのワークロードがクラッシュすると、すべてのワークロードがクラッシュします (同じGPUを共有して使っているため。例えば OOM Killed があります。)

(補足) Pascal 以降のアーキテクチャ

アーキテクチャ	プロダクト例
Hopper (2022/05)	H100
Ada Lovelace (2022/10)	GeForce 40 series
Ampere (2020/05)	GeForce 30 series, A30, A100
Turing (2018/09)	Geforce 20 series, Tesla T4
Volta (2017/12)	V100
Pascal (2016/05)	Geforce 10 series, P100

Timeslice は nvidia-smi から設定する

CUDA 11.1以降（R455+ドライバー）にて利用可能です。
4つのモードが提供されています。 (0=DEFAULT, 1=SHORT, 2=MEDIUM, 3=LONG)

$ nvidia-smi compute-policy --set-timeslice 0=DEFAULT, 1=SHORT, 2=MEDIUM, 3=LONG

このモードの違いについては明文化されたドキュメントは私の方でまだ発見できていません…
同じ悩みをもった issue もいくつかあるようでした。

• https://forums.developer.nvidia.com/t/is-there-any-document-on-set-timeslice-on-nvidia-smi/156566/2
• https://github.com/NVIDIA/k8s-device-plugin/issues/354

ちなみに私のプライベートPCの Geforce はどうなのかを確認したところ、以下の通りでした。

NGパターン (NVIDIA GeForce RTX 3060 Laptop GPU)
PS C:\Users\tozastation> nvidia-smi compute-policy --list
+-------------------------------------+
| GPU compute policies:               |
| GPU     Name           Value        |
|=====================================|
Failed to get compute policies for GPU 0 : Not Supported
Failed to retrieve compute policies for requested devices : Not Supported

GKE Time-Sharing GPU との違いはあるか

GKEのドキュメントにもあるため目に入った方もいらっしゃると思います。
ここについての違いは結論から言うとありません。

Compute Policy にて Timeslice の 設定がされています。

kubectl debug -it --image=busybox nodes/<node_name> -- ash
root@cuda-simple-5c4456dcc8-vbcxk:/# nvidia-smi compute-policy -l
+-------------------------------------+
| GPU compute policies:               |
| GPU     Name           Value        |
|=====================================|
|   0   Timeslice       Default       |
+-------------------------------------+

Time-Sharing GPU に向いているワークロードはあるか

Google Cloud のドキュメントにて少し紹介されています。

それは、GPUリクエストが少ない同系統のアプリ, バースト可能などの特徴があり、「レンダリング」, 「推論」, 「小規模な機械学習モデルのトレーニング」「ML/データ分析の開発環境」などが具体例として挙げられます。

Time-Slicing GPU を使う環境を整える

「GPUアプリケーションの実装」と「Kubernetes」の面で整理していきます。

前半では、アプリケーションが利用するGPUメモリ使用量を制御する方法と、その使用量を把握するための手法に焦点を当てます。後半では、NVIDIA/k8s-device-pluginのTime-Slicing GPU設定と、コンテナをスケジュールする際に工夫する方法について紹介します。

GPUアプリケーション

アプリが利用する GPU メモリの量を制限する

メモリ制限は Timeslice GPU には適用されません。1つのワークロードがクラッシュすると、すべてのワークロードがクラッシュします (同じGPUを共有して使っているため。例えば OOM Killed があります。)

アプリケーション側で利用するメモリ使用量を制限することで、前述したこちらの注意点を回避する方法をサンプルコードと共に紹介します。

アプローチ

• GPU Capacity に対して利用する割合を指定する or 直接サイズを指定する

検証内容

1. フレームワークに実装されている機能を使ってメモリ使用量を制限
2. ちょっとはみ出るように行列データを作成

PyTorch サンプル

• 制限に使う関数: set_per_process_memory_fraction
  • Capacity に対しての 0 ~ 1 で指定

# 実際にGPUメモリ使用量を制限しプロセスを落とすサンプルコード
import torch

mem_info = torch.cuda.mem_get_info()
print("free: " + str(round(mem_info[0]/1024**3,1)) + "GB")
print("available: " + str(round(mem_info[1]/1024**3,1)) + "GB")

memory_fraction = 0.5
print("set memory fraction: " + str(memory_fraction))
torch.cuda.set_per_process_memory_fraction(memory_fraction)
print("max use memory: " + str(round(mem_info[1]/1024**3*memory_fraction,1))+ "GB")

# torch.ones が float32 (4byte) でデータを埋めるため 欲しいサイズ/4で割ってる
tensor_size = int((3 * 1024 ** 3) / 4)
size = (tensor_size,)
tensor = torch.ones(size=size, device='cuda')

root@b219ac294b69:/workspaces/gpu # python pytorch.py 
free: 5.0GB
available: 6.0GB
set memory fraction: 0.5
max use memory: 3.0GB
Traceback (most recent call last):
  File "/workspaces/gpu/pytorch.py", line 19, in <module>
    tensor = torch.ones(size=size, device='cuda')
torch.cuda.OutOfMemoryError: CUDA out of memory. Tried to allocate 3.00 GiB (GPU 0; 6.00 GiB total capacity; 0 bytes already allocated; 5.01 GiB free; 3.00 GiB allowed; 0 bytes reserved in total by PyTorch) If reserved memory is >> allocated memory try setting max_split_size_mb to avoid fragmentation.  See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF

無事設定した閾値よりも多く使うとアプリケーションが落ちてくれました。

Tensorflow サンプル

• 制限に使う関数: tf.config.set_logical_device_configuration
  • MBの単位で直接指定

# 実際にGPUメモリ使用量を制限しプロセスを落とすサンプルコード
import tensorflow as tf
import torch

gpus = tf.config.list_physical_devices('GPU')
if gpus:
  try:
    tf.config.set_logical_device_configuration(
        gpus[0],
        [tf.config.LogicalDeviceConfiguration(memory_limit=1024)])
    logical_gpus = tf.config.list_logical_devices('GPU')
  except RuntimeError as e:
    print(e)
# Pytorch の Tensor を Numpy に変換して tf tensor に変換する
tensor_size = int((3 * 1024 ** 3) / 4)
size = (tensor_size,)
tensor = torch.ones(size=size)
np_tensor = tensor.numpy()
tf_tensor = tf.convert_to_tensor(np_tensor)

制限するGPUメモリ量をどう見極めていくか

2つご紹介します。
アプローチはどちらとも「ピーク時のGPUメモリ使用量」を取得する形になります。

① アプリ側にてメトリクスを取得する

• PyTorch
  • https://pytorch.org/docs/stable/generated/torch.cuda.memory_stats.html
• Tensorflow
  • https://www.tensorflow.org/api_docs/python/tf/config/experimental/get_memory_info

② NVIDIA/k8s-device-plugin からメトリクスを取得する

• DCGM-Exporter (DCGM: NVIDIA Datacenter GPU Manager) を k8s にインストール
  • DCGMはGPUテレメトリ (メモリ使用率・メモリクロック etc)を取得するエージェント
  • DCGM-Exporter はそのテレメトリをPrometheusで取得できるようにするコンポーネントです。
• 見たい値はこちら
  • PromQL: max_over_time(DCGM_FI_DEV_GPU_UTIL[duration])
• Grafana ダッシュボードも提供されている

Kubernetes

NVIDIA/k8s-device-plugin の調整

Time-Slicing GPU の設定方法

nvidia-device-plugin が nvidia.com/gpu というGPUリソースを kubelet を経由し登録し Kubernetes にて NVIDIA GPU を使ったアプリケーションが利用できるようになります。
Time-Slicing GPU は そのGPUリソースを幾つまで増やすかを replicas で設定します。

これにより既存のGPU数 × replicas がそのノードで利用できる量に変わります。

↓ 設定サンプル ↓ ※

apiVersion: v1
kind: ConfigMap
metadata:
  name: nvidia-device-plugin-config
data:
  any: |-
    sharing:
      timeSlicing:
        # 共有GPUのリソース名は接頭尾に.sharedをつけるようにするオプション
        renameByDefault: true
        # 共有GPUを1つのコンテナに1より多く割り当てようとした場合のバリデーションをするかどうか
        failRequestsGreaterThanOne: true
        resources:
          # resources[*].name: 利用できるリソース名 (MIGでない場合は固定)
          - name: nvidia.com/gpu
          # resources[*].replicas: 1ノードあたりにアサインできるGPUの量
          replicas: 4

マシンタイプにより利用できるリソース量を変える場合

Node Labels, NodeSelector を用いてマシンタイプごとに NVIDIA/k8s-device-plugin をデプロイする必要があります

サンプル

nodeSelector:
  workload.gpu/16GiB: "true"

GPUアプリに replicas の情報を伝える

先ほど、アプリケーション側で GPUメモリの使用量を制限する話をしました。あまり手をかけずにKubernetesで運用するには、NVIDIA/k8s-device-plugin 側で replicas を幾つにしたかアプリが知ることで柔軟に実現できそうです。

1. NVIDIA_REPLICAS などの環境変数を用意しコンテナに渡す方式
2. ノードラベルに nvidia.com/gpu.replicas がついているため、アプリから取得する

私は面倒くさがりなので 2. の方式が好きです。

まとめ

Timeslice は複数のコンテナを同じGPU上でスケジュールできるようにするための仕組みです。
NVIDIA/k8s-device-plugin は その仕組みを利用するために、共有GPU、replicas を用いて Kubernetes で利用できるよう調整します。

Timeslice を利用するには GPUアプリ (GPUメモリ使用量の制限の導入)、NVIDIA/k8s-device-plugin (各ノードが利用できるGPU数の配信) 等設定が必要です。
GPUメモリ使用量の閾値の設定は、メトリクスを用いて判断するのがおすすめでピーク時の値が参考になるかと思います。もちろん GPUアプリのタイプによっては判断がつきにくいものもあるのでそれは Timeslice で利用すべきではないと言えます。