「HAMi」でKubernetes上のGPUメモリ分離の仕組みを理解し、共有を試してみよう

はじめに

スリーシェイクのSreake事業部に所属する李(LinkedIn)です。

Kubernetesの標準的なGPUスケジューリングでは、GPUは整数単位で割り当てられ、1枚のGPUは1つのPodに占有されます。そのGPUがほとんど使われていなくても他のPodは確保できずPending状態で待機します。わずかなVRAM(Video RAM)しか使わない推論ワークロードでは、GPUの大半がアイドル状態となり、利用率が上がらずコスト効率が悪化しがちです。

この課題を解決するのが「HAMi」(Heterogeneous AI Computing Virtualization Middleware)です。KubernetesでGPUを複数のコンテナで共有できるようにするオープンソースプラットフォームです。

HAMiが提供するGPU共有は、「VRAMのリソース分離」と「GPUコアのデバイス共有(time-Slicing方式)」の2種類です。本記事では主に前者の「VRAMリソース分離」に注目し、その仕組みと動作をEKS上で確認します。

インストール周りの注意点

事前準備

本記事では、EKS環境において以下のスペックで検証を行います。

• GPU Instance Type : g4dn.xlarge、NVIDIA T4 ×1
• OS : Amazon Linux 2023.11
• Kernel : 6.18.30-61.116.amzn2023.x86_64
• Kubernetes : v1.36.1-eks-3385e9b
• Container Runtime: containerd 2.2.3

HAMiの公式ドキュメントでは、NVIDIAをcontainerdのデフォルトランタイムに設定するよう案内されていますが、EKSのGPU AMIでは設定済みのため手動設定は不要です。実際にGPUワーカーノードに接続し、状態を確認します。

# default runtime 
$ containerd config dump 2>/dev/null | grep 'default_runtime_name'
default_runtime_name = 'nvidia'

HAMiのスケジューリング対象とするGPUノードには、gpu=onラベルを付与します。このラベルが付いていないノードは管理対象になりません。今回はrole=gpuラベルが付いたノードにgpu=onを付与します。

# role ラベルを確認
$ kubectl get nodes -L role
NAME                                             STATUS   ROLES    AGE    VERSION               ROLE
ip-10-0-45-164.ap-northeast-1.compute.internal   Ready    <none>   127m   v1.36.1-eks-3385e9b   system
ip-10-0-45-70.ap-northeast-1.compute.internal    Ready    <none>   118m   v1.36.1-eks-3385e9b   gpu


# role=gpuノードへgpu=onを付与
$ kubectl label nodes -l role=gpu gpu=on

# role=gpu というラベルが付いたノードだけ表示
$ kubectl get nodes -l role=gpu
NAME                                            STATUS   ROLES    AGE    VERSION
ip-10-0-45-70.ap-northeast-1.compute.internal   Ready    <none>   119m   v1.36.1-eks-3385e9b

# gpu ラベルの値を追加し、全ノードを表示
$ kubectl get nodes -L gpu
NAME                                             STATUS   ROLES    AGE    VERSION               GPU
ip-10-0-45-164.ap-northeast-1.compute.internal   Ready    <none>   127m   v1.36.1-eks-3385e9b   
ip-10-0-45-70.ap-northeast-1.compute.internal    Ready    <none>   119m   v1.36.1-eks-3385e9b   on

Helmを用いたHAMiインストール

HelmでHAMiをインストールします。検証環境ではGPUノードにtaintが付与されているため、taintを確認してhami-device-pluginにtolerationを設定します。設定しないとdevice pluginのPodがGPUノードにスケジュールされず、GPUがリソースとして登録されません。

$ helm repo add hami-charts https://project-hami.github.io/HAMi
$ helm repo update

# hami-device-plugin をインストールするため、Taint を確認
$ kubectl get nodes -l role=gpu -o jsonpath='{range .items[*]}{.metadata.name}{"\t"}{.spec.taints}{"\n"}{end}'
ip-10-0-44-102.ap-northeast-1.compute.internal   [{"effect":"NoSchedule","key":"nvidia.com/gpu","value":"present"}]
 
# hami-device-pluginにtolerationを設定
$ helm install hami hami-charts/hami \
  -n kube-system \
  --set scheduler.kubeScheduler.imageTag="v1.36.1" \
  --set 'devicePlugin.tolerations[0].key=nvidia.com/gpu' \
  --set 'devicePlugin.tolerations[0].operator=Exists' \
  --set 'devicePlugin.tolerations[0].effect=NoSchedule'

インストール後、GPUノードに1枚のGPUデバイス(NVIDIA-Tesla T4 16GB)が登録されました。

$ kubectl get nodes -l role=gpu -o json | \
  jq -r '.items[].metadata.annotations
         | to_entries[] | select(.key|test("hami|4pd|gpu"))
         | "\(.key) = \(.value)"'

hami.io/node-nvidia-register = [{"id":"GPU-32f41276-93b5-b080-ede7-533489f70b70","count":10,"devmem":15360,"devcore":100,"type":"NVIDIA-Tesla T4","mode":"hami-core","health":true,"devicepairscore":{}}]

HAMiコンポーネントとCUDA APIフックの話

HAMiのコンポーネント

HAMiは主に「hami-scheduler」と「hami-device-plugin」で構成されます。

$ kubectl -n kube-system get pods | grep -i hami
hami-device-plugin-4k2sn          2/2     Running   0          53s
hami-scheduler-6d5886db4d-5ql6p   2/2     Running   0          53s

hami-schedulerは標準スケジューラを拡張します。標準スケジューラはnvidia.com/gpumem(VRAM)やnvidia.com/gpucores(コア)といったHAMi独自のリソース要求を理解できないため、hami-schedulerがこれらを解釈し、各ノードのGPUの空き状況をもとに配置先を決定します。そして、どのGPUにどれだけ割り当てたかをPodのアノテーションに書き込みます。

hami-device-pluginは、kubeletへのGPUリソース登録(物理GPU 1枚を論理的なN枚に拡張し、kube-schedulerにN枚あると認識させる)と、ノードアノテーションへのデバイス仕様(VRAM容量、UUID、演算能力など)の記録を行います。さらにPod生成時には、hami-schedulerが書き込んだ割り当て情報をもとに、コンテナへlibvgpu.soとVRAM上限値を注入します。

このlibvgpu.soがVRAM分離の中核です。注入されたパスはコンテナ内の/etc/ld.so.preloadに記載され、コンテナ起動時に優先的にロードされます。

$ kubectl exec -it hami-test -- cat /etc/ld.so.preload
/usr/local/vgpu/libvgpu.so

CUDA APIフックによるVRAM分離の仕組み

libvgpu.soには、CUDAのメモリ確保系APIの呼び出しを横取りする関数が定義されています。

# libvgpu.so が CUDA のメモリ確保系 APIを横取り
$ nm -D ./libvgpu.so | grep -iE 'cuMemAlloc|cudaMalloc|nvml'
0000000000030ada T cuMemAllocAsync
0000000000031953 T cuMemAllocFromPoolAsync
0000000000029f43 T cuMemAllocHost_v2
000000000002a246 T cuMemAllocManaged
000000000002a520 T cuMemAllocPitch_v2
0000000000029db0 T cuMemAlloc_v2

GPU上のアプリケーションがVRAMを確保する際は、最終的にこれらのCUDA Driver APIを呼び出します。

例えば、PyTorchやTensorFlowはフレームワーク内部でこれらのAPIを呼ぶため、その呼び出しがlibvgpu.soに横取りされます。この仕組みはフレームワークに依存せず、CUDA Driver APIを経由してメモリを確保するアプリケーションであれば同様に機能します。

横取りされた確保要求は、環境変数CUDA_DEVICE_MEMORY_LIMITに設定された上限値と照合されます。上限内であれば本来のCUDA APIに処理を渡して確保を許可し、超過する場合はエラーを返して拒否します。これにより、コンテナから見えるVRAM総量が制限され、上限を超える要求はCUDAのOOM(Out of Memory)として失敗します。

Podのデプロイ方法

HAMiでVRAMを分離するには、Podのマニフェストにnvidia.com/gpumemを定義します。

ここでnvidia.com/gpuとnvidia.com/gpumemの違いに注意が必要です。nvidia.com/gpuは整数値でしか指定できず、これは「使用するGPUの枚数」を表します。一方、nvidia.com/gpumemは、その1枚のうち何MiBを使うかを表します。

したがって、nvidia.com/gpumemを指定しない場合、PodはGPU 1枚をそのまま占有します。仮に1つのGPUノードしかなかった場合、2つのPodがそれぞれnvidia.com/gpu: 1のみを指定すると、片方はGPUを確保できずPendingになります。

なお、HAMiを使う場合でも、GPUノードが複数あり、各ノードにPodを1つずつ配置したい場合は、Pod anti-affinity(topologyKey: kubernetes.io/hostname)を指定することで同じノードに複数のPodが集中することを防ぎ、各ノードへ分散して配置できます。

$ kubectl apply -f hami-sample.yaml

【例】VRAMを分割して利用するPodのマニフェスト

# hami-sample.yaml
apiVersion: v1
kind: Pod
metadata:
  name: hami-test
spec:
  containers:
  - name: cuda
    image: nvidia/cuda:12.4.0-base-ubuntu22.04
    command: ["sleep", "3600"]
    resources:
      limits:
        nvidia.com/gpu: 1
        nvidia.com/gpumem: 2000      #　VRAMの分離(2GiB)
  tolerations:
  - key: nvidia.com/gpu
    operator: Exists
    effect: NoSchedule

デバイス共有の検証

ここからは、HAMiの「デバイス共有」機能を検証していきます。1つのGPUを2つのPodで共有できることを確認します。

以下のマニフェストでは、検証用に2つのPod(core-a、core-b)を生成、それぞれ4GiBのVRAMを要求しています。

kubectl apply -f hami-share-device.yaml

# hami-share-device.yaml
apiVersion: v1
kind: Pod
metadata:
  name: core-a
spec:
  restartPolicy: Never
  containers:
  - name: torch
    image: pytorch/pytorch:2.4.1-cuda12.4-cudnn9-runtime
    command: ["sleep", "infinity"]
    resources:
      limits:
        nvidia.com/gpu: 1
        nvidia.com/gpumem: 4000 # VRAM 4GiB
---
apiVersion: v1
kind: Pod
metadata:
  name: core-b
spec:
  restartPolicy: Never
  containers:
  - name: torch
    image: pytorch/pytorch:2.4.1-cuda12.4-cudnn9-runtime
    command: ["sleep", "infinity"]
    resources:
      limits:
        nvidia.com/gpu: 1
        nvidia.com/gpumem: 4000 # VRAM 4GiB

nvidia-smiコマンドで各Podが見ているGPUのUUIDを確認すると、2つのPodのUUIDが同一であり、1枚の物理GPUを共有していることが分かります。

# 2つの UUID が同一 → GPU共有を確認
$ kubectl exec core-a -- nvidia-smi --query-gpu=uuid --format=csv,noheader
GPU-f7994ebc-82e7-67a4-6aba-df11a880d49d

$ kubectl exec core-b -- nvidia-smi --query-gpu=uuid --format=csv,noheader
GPU-f7994ebc-82e7-67a4-6aba-df11a880d49d

デバイスリソース分離の検証

続いて、HAMiの「デバイスリソース分離」機能を検証します。共有している2つのPodが、それぞれ指定したVRAM上限を超えられないことを確認します。

まず、Pod(core-a)に割り当てられたVRAMを確認します。要求どおり4000 MiBに制限されています。

$ kubectl exec core-a -- nvidia-smi --query-gpu=memory.total --format=csv
memory.total [MiB]4000 MiB

次に、要求サイズ(GiB)を引数に取るサンプルPythonコードを作成します。try_alloc()の引数にVRAMのサイズを指定し、上限内と上限超過の両方を試します。

$ kubectl exec core-a -- bash -c 'cat > /tmp/test.py << "EOF"
import torch

def try_alloc(gb):
    try:
        n = int(gb * 1024**3 // 4)  # float32 = 4 bytes
        x = torch.zeros(n, dtype=torch.float32, device="cuda")
        mib = x.element_size() * x.nelement() // 1024**2
        print(f"[OK]   {gb}GB allocated ({mib} MiB)")
        del x
        torch.cuda.empty_cache()
    except RuntimeError as e:
        print(f"[FAIL] {gb}GB -> {type(e).__name__}: {str(e)[:120]}")

print("visible total:", torch.cuda.get_device_properties(0).total_memory // 1024**2, "MiB")
try_alloc(2) # 2GiBを要求
try_alloc(6) # 6GiBを要求
EOF'

サンプルPythonコードを実行します。

$ kubectl exec core-a -- python /tmp/test.py

visible total: 4000 MiB
[OK] 2GB allocated (2048 MiB)
[FAIL] 6GB -> OutOfMemoryError: CUDA out of memory. Tried to allocate 6.00 GiB. GPU 0 has a total capacity of 3.91 GiB of which 3.72 GiB is free.

結果から、上限内の2GiB要求は成功し、上限を超える6GiB要求はOOMで失敗することが分かります。libvgpu.soによってVRAMがハードに分離されていることが確認できました。

おわりに

本記事では、EKS環境でHAMiを用いたGPUの「デバイス共有」と「デバイスリソース分離」を検証しました。

HAMiの大きな強みは、特定モデルでしか使えないMIG(Multi-Instance GPU)とは異なり、MIG非対応のGPUでもVRAMの隔離を実現できる点です。間欠的にしかGPUを使わないCIタスクや、数GB程度のVRAMしか必要としない小規模言語モデルの推論サービスなどで特に効果を発揮します。

一方、GPUコアはVRAMのようなハードな分離ではなくtime-slicing方式のため、複数Podが同時にGPUを使用すると演算性能が互いに影響し合う可能性があります。厳密な性能保証にはMIGが、GPUをフルに使い切らないワークロードにはHAMiが適していると言えるでしょう。

【参考】

• https://project-hami.io/docs
• https://jimmysong.io/blog/kubernetes-gpu-control-plane-hami-v29-ai-infra/