コンテンツにスキップ

kfp local の実行時に GPU を使う

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

環境

  • kfp v2.7.0

Docker の runtime を nvidia に変更する

docker CLI で GPU を使うにはdocker run --gpus all <IMAGE>のように実行することで GPU を使うことはできる。しかし、kfp local の実行時に、 CLI 実行のようなオプションを指定できない。 設定ファイルを編集し、Docker の runtime を nvidia に変更することで kfp local の環境で GPU が利用できる。

/etc/docker/daemon.json
{
  "default-runtime": "nvidia",
  "default-shm-size": "16G",
  "runtimes": {
    "nvidia": {
      "path": "nvidia-container-runtime",
      "runtimeArgs": []
    }
  }
}

次のコマンドのように docker を再起動すると変更が反映され、kfp local の実行時に GPU を利用できる。

sudo systemctl restart docker