title: Running custom Docker container with GPU using Vertex AI Pipelines - how to install NVIDIA driver?
body: <p>I need to run a custom Docker container with GPU support using Vertex AI Pipelines, and I'm not seeing a clear way to do that. This requires several components:</p>
  <ol>
  <li>Applications (pytorch in my case)</li>
  <li>CUDA toolkit</li>
  <li>CUDA driver</li>
  <li>NVIDIA GPUs</li>
  </ol>
  <p>I can use a <a href="https://github.com/NVIDIA/nvidia-docker" rel="nofollow noreferrer">NVIDIA Docker</a> base image for #1 and #2, and a GCP accelerator for #4, but how do I install the CUDA driver in a Vertex AI pipeline? There's documentation on how to install NVIDIA drivers <a href="https://cloud.google.com/compute/docs/gpus/install-drivers-gpu#no-secure-boot" rel="nofollow noreferrer">on GCE instances</a> and <a href="https://cloud.google.com/kubernetes-engine/docs/how-to/gpus#installing_drivers" rel="nofollow noreferrer">GKE nodes</a>, but nothing for Vertex AI.</p>
  <p>One option could be to <a href="https://cloud.google.com/deep-learning-containers/docs/derivative-container" rel="nofollow noreferrer">create a derivative container based on a GCP Deep Learning Container</a>, but then I have to use a GCP container and don't have as much control over the environment.</p>