😎

WSLのkubernetesでGPUを使う

2023/03/25に公開

はじめに

いろいろな事情(会社PCの標準OSなどなど)で、Windowsを使わないといけないことは多いと思います。Windows嫌いではないのですが、何かとLinuxの方が動かしやすいです。なのでWSLを常用しています。
今回はkubernetesの勉強をしてみたくなったので、WSLにmicrok8sで環境構築をしてGPUも扱えるようにしました。

環境

Windows 11
Nvidia GeForce RTX 3080 / RTX 3050 Ti Laptop
WSL Version: 1.1.3.0
Ubuntu 22.04 LTS
microk8s v1.26

まずは結論

長くなりそうなのでまず結論を先に書いておきます。
Nvidia公式のNvidia-device-pluginをWSLでも動くよう修正して動かすことができました。
使い方はREADMDにも書きました。
https://github.com/rrrrrrryo/k8s-device-plugin-wsl-gpu

microk8s

microk8sはCanonical社が開発している軽量kubernetes環境です。
オリジナルのkubernetes環境構築ツールであるkubeadmのほかに、軽量な環境であればminikube、k3s、kindなんかが有名だと思います。ここらへんの比較は、microk8s公式ページで比較されています。
https://microk8s.io/compare
snapコマンドで簡単に入れたり消したりできるmicrok8sを選びました。
また、microk8sはadd-onとして色々な環境構築を手助けする機能があるので簡単かな思います。
(例えば、dashboarddnsingressなど)

GPUを使うために

gpuというadd-onがあり、本来はこの機能を使えばmicrok8sのkubnernetesでgpuが簡単に使えるようになるはずでした。
しかし、WSL環境ではそんな簡単にいきません。また、使っているGPUによってはmicrok8sのgpu add-onではうまくいかないと思います。microk8sでgpuを使うための公式ドキュメントは以下の通りです。
https://microk8s.io/docs/addon-gpu
このアドオンは、v1.21以降からNvidia GPU Operatorというdevice pluginを使っているようです。しかし、このOperatorのGPUデバイスサポートは公式の通り、AやTシリーズなどです。
https://docs.nvidia.com/datacenter/cloud-native/gpu-operator/platform-support.html#supported-nvidia-gpus-and-systems
このため、私の環境では素直にアドオンを使ってGPUを制御できません。
アドオンを使わずにNvidia-device-pluginで、GPUを使えるようにすれば良いかなと思います。
https://github.com/NVIDIA/k8s-device-plugin
このdevice-pluginを使えば簡単に動くかなと思っていましたが、GPUを割り当てたPodはPendingのまま立ち上がりません。
issueを見ると、WSLはGPUを/dev/dxgで仮想化してあり、通常の/dev/nvidia*で管理されていないためということでした。
https://github.com/NVIDIA/k8s-device-plugin/issues/332

nvidia-device-pluginを修正

上のissueの通り、そのままだとWSLでは動きません。ということでnvidia-device-pluginを修正します。WSLの/dev/dxgの仮想化とNvidiaのnvmlが素晴らしいので大きな修正は必要ありませんでした。/dev/nvidia*でデバイスファイルを探している箇所を、/dev/dxgがあればWSLとしてデバイスを扱います。
修正箇所の一例を以下に示します。

internal/rm/nvml_devices.go
@@ -68,11 +68,17 @@ func (d nvmlMigDevice) GetUUID() (string, error) {

// GetPaths returns the paths for a GPU device
    func (d nvmlDevice) GetPaths() ([]string, error) {
-	    minor, ret := d.GetMinorNumber()
-	    if ret != nvml.SUCCESS {
-		    return nil, fmt.Errorf("error getting GPU device minor number: %v", ret)
+	    path := ""
+	    if _, err := os.Stat("/dev/dxg"); os.IsNotExist(err){
+		    minor, ret := d.GetMinorNumber()
+		    if ret != nvml.SUCCESS {
+			    return nil, fmt.Errorf("error getting GPU device minor number: %v", ret)
+		    }
+		    path = fmt.Sprintf("/dev/nvidia%d", minor)
+
+	    } else {
+		    path = "/dev/dxg"
+	    }
-	    path := fmt.Sprintf("/dev/nvidia%d", minor)

	    return []string{path}, nil
}

ということで、device-pluginを修正しこれを動かすpodを用意します。

手順

以下は、WSLでnvidia-docker2が動く状態であることを前提とします。

  1. microk8sをインストールする
    $ sudo snap install microk8s --classic
    
  2. microk8sのcontainerdのruntimeをnvidia-container-runtimeを使えるように設定する
    microk8sの公式ドキュメントの通りですが、私の環境では少し修正が必要でした。
    $ echo '
            [plugins."io.containerd.grpc.v1.cri".containerd.runtimes.nvidia]
            runtime_type = "io.containerd.runc.v2"
    
            [plugins."io.containerd.grpc.v1.cri".containerd.runtimes.nvidia.options]
                BinaryName = "/usr/bin/nvidia-container-runtime"
    ' | sudo tee -a /var/snap/microk8s/current/args/containerd-template.toml
    
    これに加え、containerd-template.tomldefault_runtime_name"nvidia"に変更します。
    その後microk8sを再起動します。
    $ sudo snap restart microk8s
    
  3. k8s-device-plugin-wsl-gpuをbuildしてmicrok8sに取り込む
    k8s-device-plugin-wsl-gpuをクローンしておきます。
    $ cd <k8s-device-plugin-wsl-gpu repo>
    $ docker build \
      -t k8s-device-plugin-wsl:devel \
      -f deployments/container/Dockerfile.ubuntu \
      .
    $ docker save k8s-device-plugin-wsl:devel > k8s-device-plugin-wsl.tar
    $ microk8s ctr images import k8s-device-plugin-wsl.tar
    
  4. device-pluginのpodを動かす
    $ kubectl apply -f nvidia-device-plugin.yml
    

以上でkubernetesでGPUリソースを管理する準備ができました。

試しに動かしてみます。

$ cat <<EOF | kubectl apply -f -
apiVersion: v1
kind: Pod
metadata:
  name: gpu-pod
spec:
  restartPolicy: Never
  containers:
    - name: cuda-container
      image: nvcr.io/nvidia/k8s/cuda-sample:vectoradd-cuda10.2
      resources:
        limits:
          nvidia.com/gpu: 1 # requesting 1 GPU
  tolerations:
  - key: nvidia.com/gpu
    operator: Exists
    effect: NoSchedule
EOF

$ kubectl logs pod/gpu-pod
[Vector addition of 50000 elements]
Copy input data from the host memory to the CUDA device
CUDA kernel launch with 196 blocks of 256 threads
Copy output data from the CUDA device to the host memory
Test PASSED
Done

おわり

WSLのkubernetesでGPUを管理できるようになりました。
WSLやNVMLやnvidia-container-toolsがとてもよくできていたため、すんなり動きました。
注意としては、WSLではGPUを分割して管理できないため、現時点は全てを使う割り当てしかできないと思います。
nvidia-dockerで--gpus allのみサポートされているのと同様です。

GitHubで編集を提案

Discussion