kfp local の実行時に GPU を使う¶

kfp(kubeflow pipeline)には local 環境で実行する機能がある。 local 実行は Docker 環境か subprocess 環境で実行される。今回は Docker 環境で実行する場合に GPU を使う方法を書き残す。

環境¶

• kfp v2.7.0

Docker の runtime を nvidia に変更する¶

デフォルトの Docker runtime であってもdocker run --gpus allで GPU を使うことはできるが、kfp local の実行時にはこのオプションを指定できない。 次のように Docker の runtime を nvidia に変更することで GPU を使えるようにする。

1
2
3
4
5
6
7
8
9

{
  "default-runtime": "nvidia",
  "runtimes": {
    "nvidia": {
      "path": "nvidia-container-runtime",
      "runtimeArgs": []
    }
  }
}

docker を再起動すると変更が反映され、kfp local の実行時に GPU を使うことができる。

1

sudo systemctl restart docker